autoregressive_transformer.py

import torch
from torch import nn
from tqdm import tqdm


def subsequent_mask(size):
    return torch.triu(torch.full((size, size), float('-inf')), diagonal=1)


if __name__ == "__main__":
    d_model = 512

    transformer_encoder = nn.TransformerEncoderLayer(d_model=d_model, nhead=8)

    # Initial state
    x = torch.zeros(1, 1, d_model)

    # Loop, concatenate
    for i in tqdm(range(500)):
        sequence_length = x.size(0)
        src_mask = subsequent_mask(sequence_length)

        # in iteration i - 1, sequence of length i attends to sequence of length i
        # but we only want the last element to attend to the sequence
        output = transformer_encoder(x, src_mask=src_mask)
        last_output = output[-1:]
        x = torch.cat([x, last_output], dim=0)

    # Computation is cubic in sequence length instead of quadratic