I am using a Transformer for next frame prediction. Every frame has been previously encoded into a 1D latent vector using a VAE (the encoding-decoding from the VAE is very good).

I have divided the videos into chunks of 24 frames, and using the VAE I have created sequences like 24,100 (where 24 is the sequence length, and 100 the latent vectors from the VAE). These sequences are inputted into a small Transformer (GPT-like, only the decoder with casual attention). Because it is batched, the actual input is 64,24,100. Let’s call one of those “source”. I input to the transformer “source[:,:-1]” and the target is “source[:,1:]”. The loss function is MSE. When I check the outputs from the transformer, it does an excellent prediction in all the sequence positions (0,23), except in the last one (24). The last one is OK’ish. Not random, but not amazing. As an example, all the other positions get a MSE loss of around 0.006, while the last one gets an MSE loss of 0.15.

These would all indicate that I am not using casual masking – since it does well in all the sequence entries where it can see the future, but it struggles in the last position where it can’t. But I am using casual attention, so none of the entries should see the future:

class BaseAttention(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super().__init__()
    self.mha = tf.keras.layers.MultiHeadAttention(**kwargs)
    self.layernorm = tf.keras.layers.LayerNormalization()
    self.add = tf.keras.layers.Add()

class CausalSelfAttention(BaseAttention):
  def call(self, x):
    attn_output = self.mha(
        query=x,
        value=x,
        key=x,
        use_causal_mask = True)
    x = self.add([x, attn_output])
    x = self.layernorm(x)
    return x

And the transformer blocks using this:

class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.3):
        super(TransformerBlock, self).__init__()
        self.att = CausalSelfAttention(
            num_heads=num_heads, key_dim=embed_dim, dropout=0.1)
        self.ffn = keras.Sequential(
            [layers.Dense(ff_dim, activation="relu"),
             layers.Dense(embed_dim), ]
        )
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout2 = layers.Dropout(rate)


    def call(self, inputs):
        attention_output = self.att(inputs)
        ffn_output = self.ffn(attention_output)
        ffn_output = self.dropout2(ffn_output)
        return self.layernorm2(attention_output + ffn_output)

Previous to this I was calculating the masks myself and sending them when creating the MHA instead of using “use_casual_mask=True”. but the results were the same.

Here you can also see my train_step, where you dan see I swift the window for inputs and predictions:

    def train_step(self, batch):
        """Processes one batch inside model.fit()."""
        source = batch[0]/4.  # [0] are the sequences of latent vectors
        source_input = source[:, :-1]
        source_target = source[:, 1:]
        with tf.GradientTape() as tape:
            preds = self(source_input)
            loss = self.compiled_loss(source_target, preds)*4.
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        self.loss_metric.update_state(loss)

        return {"loss": self.loss_metric.result()}

Here you can see my embedding and positional embedding:

class PositionEmbedding(layers.Layer):
    def __init__(self, maxlen, embed_dim):
        super(PositionEmbedding, self).__init__()
        # self.latent_emb = layers.Dense(embed_dim)
        self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)
        self.latent_emb = keras.Sequential(
            [layers.Conv2D(1, (5,5), padding="same", activation="relu"),
             layers.Reshape([maxlen, 100]),
             layers.Conv1D(embed_dim, 3, padding="same"), ]
        )

    def call(self, x):

        maxlen = tf.shape(x)[1]
        positions = tf.range(start=0, limit=maxlen, delta=1)
        positions = self.pos_emb(positions)

        x = self.latent_emb(tf.expand_dims(x,-1))
        # x = self.latent_emb(x)
        return x + positions

As you can see I have tried two things (if you check the comments). First I tried using a Dense layer as embedding, and that’s it. Then I tried using Conv2d and Conv1d (as the one uncommented now). The results were the same – using Convs it was a bit better, but not by much. The difference in predictions between [0,-1] and [-1] remains.

I am confused why there’s such a big difference, because using Casual Masks, it shouldn’t be like this. I have also check to see if the problem is with the padding doing the Conv, but the first prediction (row) seems to be unaffected. First and last columns also seem to be OK. The only problem is the last row (the new prediction). I tried different padding techniques, but the problem remained. If the Conv padding was the problem you would expect it all around the data, not just in the last row.

Overall I am confused, because it all seems to indicate it is ignoring the casual mask, but to the best of my knowledge it is using it.