Transformer model is very slow and doesn't predict well

Question

I created my first transformer model, after having worked so far with LSTMs. I created it for multivariate time series predictions - I have 10 different meteorological features (temperature, humidity, windspeed, pollution concentration a.o.) and with them I am trying to predict time sequences (24 consecutive values/hours) of air pollution. So my input has the shape X.shape = (75575, 168, 10) - 75575 time sequences, each sequence contains 168 hourly entries/vectors and each vector contains 10 meteo features. My output has the shape y.shape = (75575, 24) - 75575 sequences each containing 24 consecutive hourly values of the air pollution concentration.

I took as a model an example from the official keras site. It is created for classification problems, I only took out the softmax activation and in the last dense layer I set the number of neurons to 24 and I hoped it would work. It runs and trains, but it does worse predictions than the LSTMs I have used on the same problem and more importantly - it is very slow - 4 min/epoch. Below I attach the model and I would like to know:

I) Have I done something wrong in the model? can the accuracy or speed be improved? Are there maybe some other parts of the code I need to change for it to work on regression, not classification problems?

II) Also, can a transformer at all work on multivariate problems of my kind (10 features input, 1 feature output) or do transformers only work on univariate problems? Tnx

def build_transformer_model(input_shape, head_size, num_heads, ff_dim, num_transformer_blocks, mlp_units, dropout=0, mlp_dropout=0):

    inputs = keras.Input(shape=input_shape)
    x = inputs
    for _ in range(num_transformer_blocks):

        # Normalization and Attention
        x = layers.LayerNormalization(epsilon=1e-6)(x)
        x = layers.MultiHeadAttention(
            key_dim=head_size, num_heads=num_heads, dropout=dropout
        )(x, x)
        x = layers.Dropout(dropout)(x)
        res = x + inputs

        # Feed Forward Part
        x = layers.LayerNormalization(epsilon=1e-6)(res)
        x = layers.Conv1D(filters=ff_dim, kernel_size=1, activation="relu")(x)
        x = layers.Dropout(dropout)(x)
        x = layers.Conv1D(filters=inputs.shape[-1], kernel_size=1)(x)
        x = x + res

    x = layers.GlobalAveragePooling1D(data_format="channels_first")(x)
    for dim in mlp_units:
        x = layers.Dense(dim, activation="relu")(x)
        x = layers.Dropout(mlp_dropout)(x)
    x = layers.Dense(24)(x)
    return keras.Model(inputs, x)

model_tr = build_transformer_model(input_shape=(window_size, X_train.shape[2]), head_size=256, num_heads=4, ff_dim=4, num_transformer_blocks=4, mlp_units=[128], mlp_dropout=0.4, dropout=0.25)
model_tr.compile(loss="mse",optimizer='adam') 
m_tr_history = model_tr.fit(x=X_train, y=y_train, validation_split=0.25, batch_size=64, epochs=10, callbacks=[modelsave_cb])

Answer 1

Edit

I just noticed that the model you are referring to is built very differently than the transformer from Attention is All You Need since it only uses one half of the architecture. Thus my answer below is not be complete. I thus have to add the following: (The final two paragraphs still apply as they are, though)

The Keras model is quite weird, and while it is a time-series model, it's not a regression one. Thus my general description regarding how transformers work still apply (that is, you only get one prediction at a time), since their model is a many-to-one classification model. So, you have adjust your model accordingly.

The main thing you can ignore from my original answer is my emphasis on decoders; Keras only use encoders for their model.

---

Original answer

Transformers only output one prediction at a time, not twenty-four. Lets break a transformer during runtime step by step.

In all of the steps, the encoder input will be the same since your sequences are not that long. The decoder gets the output as its input, which is done in the following manner:

Step one: The decoders input is only one token which stands for "start of sequence" and padding. $(start,0,0,..,0,0)$

Step two: The decoder input is the start-token and the prediction from step one + padding $(start,y_1,0,0,..,0,0)$

Step three: The decoder input is the decoders input in step two and the prediction from step two + padding $(start,y_1,y_2,0,0,..,0,0)$

And, so on.. If you want 24 outputs, the transformer will have to run 24 times. The only reason the original transformer model has multiple outputs is that it needs to make one prediction for each word token (and, pick the most likely one). Training will thus have to reflect this.

However, I am quite skeptical that you will have any significant benefit from using transformers over LSTMs. The main benefit of transformers, or any other attention based model, is that they allow for longer-term dependencies than LSTMs. According to a Stanford study, the point at which LSTMs are no longer efficient is a sequence of around 200 tokens. Of course, that study is specific to language models, but I bet it won't be that much different in your case; Your sequences are quite short.

Anyhow, the output layer for a regression task generally use linear activation functions. If all your outputs are positive, that won't make much of a difference.