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. I runs and trains, but it doesn't do a better job 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

The bit of this code that surprised me matches the tutorial you linked to, but they seem to get very poor performance. Their val_loss at the (best) epoch 105 is 0.3232; whereas in the RNN tutorial they had a best val_loss of 0.0895 (epoch 376).

The bit that surprised me was using Conv1D(filters=4, kernel_size=1) as the feed-forward network. I looked a couple of times for a link to a paper where this had come from, or some justification/explanation, but didn't spot anything.

More normal, for a Transformer, is a simple fully-connected network. E.g. this keras tutorial does it as:

  self.ffn = keras.Sequential(
      [layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim),]
      )

which is then used as:

  out1 = self.layernorm1(inputs + attn_output)
  ffn_output = self.ffn(out1)
  ffn_output = self.dropout2(ffn_output, training=training)

Typically ff_dim is 4x the size of embed_dim; I've done experiments with trying a smaller multiplier, but I think you want it to be at least 2x.

BTW, I'm having trouble working out what your embed_dim is; but maybe it is only 10? In which case a head_dim of 256 is inappropriate. (E.g. if embed_dim was 256, you might have 4 heads each of size 64.)

Taking a step back, the way you've described your data sounds like it can be rephrased in NLP terms as you have a batch of 75575 sentences, each with exactly 168 tokens. In NLP a token is a string of unicode characters, and the embedding layer will turn them into e.g. 256 floating point values, usually each randomly initialized to be roughly -1 to +1, but then they can be learned; whereas you have a set of 10 floating-point (?) features, that haven't been normalized to that kind of range (?).

In summary, I'd first follow a tutorial that shows it managed an improvement over an RNN; and second I'd work on the data representation.