Transformers using PyTorch : Worklog Part 2

Avishek Sen Gupta on 14 January 2023

---

We continue looking at the Transformer architecture from where we left from Part 1. When we’d stopped, we’d set up the Encoder stack, but had stopped short of adding positional encoding, and starting work on the Decoder stack. In this post, we will focus on setting up the training cycle.

Specifically, we will cover:

• Positional Encoding
• Decoder stack, including the masked multi-head attention mechanism
• Set up the basic training regime via Teacher Forcing

We will also lay out the dimensional analysis a little more clearly, and add necessary unit tests to verify intended functionality. The code is available here.

Positional Encoding

You can see the code for visualising the positional encoding here. Both images below show the encoding map at different levels of zoom.

The code in the main Transformer implementation which implements the positional embedding is shown below.

    # The encoder output is injected directly into the sublayer of every Decoder. To build up the chain of Decoders
    # in PyTorch, so that we can put the full stack inside a Sequential block, we simply inject the encoder output
    # to the root Decoder, and have it output the encoder output (together with the actual Decoder output) as part of
    # the Decoder's actual output to make it easy for the next Decoder in the stack to consume the Encoder and Decoder
    # outputs
    def forward(self, input):
        encoder_output, previous_stage_output = input
        masked_mh_output = self.masked_multiheaded_attention_layer(
            self.masked_qkv_source.forward(previous_stage_output))
        input_qkv = self.unmasked_qkv_source.forward((encoder_output, masked_mh_output))
        mh_output = self.multiheaded_attention_layer(input_qkv)
        # Adds the residual connection to the output of the attention layer
        layer_normed_multihead_output = self.layer_norm(mh_output + previous_stage_output)
        ffnn_outputs = torch.stack(
            list(map(lambda attention_vector: self.feedforward_layer(attention_vector), layer_normed_multihead_output)))
        layer_normed_ffnn_output = self.layer_norm(ffnn_outputs + layer_normed_multihead_output)
        return (encoder_output, layer_normed_ffnn_output)

Data Flow

The diagram below (you’ll need to zoom in) shows the data flow for a single Encoder/Decoder, with 8 attention blocks per multihead attention layer. \(n\) represents the number of words passed into the Encoder. \(m\) represents the number of words passed into the Decoder. \(V\) represents the length of the full vocabulary.

The dimensions of the data at each stage are depicted to facilitate understanding.

graph LR; subgraph Encoder encoder_src[Source Text]--nx512-->pos_encoding[Positional Encoding]; pos_encoding--nx512-->qkv_encoder[QKV Layer] qkv_encoder--Q=nx64-->multihead_attn_1[Attention 1] qkv_encoder--K=nx64-->multihead_attn_1 qkv_encoder--V=nx64-->multihead_attn_1 qkv_encoder--Q=nx64-->multihead_attn_2[Attention 2] qkv_encoder--K=nx64-->multihead_attn_2 qkv_encoder--V=nx64-->multihead_attn_2 qkv_encoder--Q=nx64-->multihead_attn_3[Attention 3] qkv_encoder--K=nx64-->multihead_attn_3 qkv_encoder--V=nx64-->multihead_attn_3 qkv_encoder--Q=nx64-->multihead_attn_4[Attention 4] qkv_encoder--K=nx64-->multihead_attn_4 qkv_encoder--V=nx64-->multihead_attn_4 qkv_encoder--Q=nx64-->multihead_attn_5[Attention 5] qkv_encoder--K=nx64-->multihead_attn_5 qkv_encoder--V=nx64-->multihead_attn_5 qkv_encoder--Q=nx64-->multihead_attn_6[Attention 6] qkv_encoder--K=nx64-->multihead_attn_6 qkv_encoder--V=nx64-->multihead_attn_6 qkv_encoder--Q=nx64-->multihead_attn_7[Attention 7] qkv_encoder--K=nx64-->multihead_attn_7 qkv_encoder--V=nx64-->multihead_attn_7 qkv_encoder--Q=nx64-->multihead_attn_8[Attention 8] qkv_encoder--K=nx64-->multihead_attn_8 qkv_encoder--V=nx64-->multihead_attn_8 subgraph EncoderMultiheadAttention[Encoder Multihead Attention] multihead_attn_1--nx64-->concat((Concatenate)) multihead_attn_2--nx64-->concat multihead_attn_3--nx64-->concat multihead_attn_4--nx64-->concat multihead_attn_5--nx64-->concat multihead_attn_6--nx64-->concat multihead_attn_7--nx64-->concat multihead_attn_8--nx64-->concat end concat--nx512-->linear_reproject[Linear Reprojection] linear_reproject--1x512-->ffnn_encoder_1[FFNN 1] linear_reproject--1x512-->ffnn_encoder_2[FFNN 2] linear_reproject--1x512-->ffnn_encoder_t[FFNN x] linear_reproject--1x512-->ffnn_encoder_n[FFNN n] subgraph FfnnEncoder[Feed Forward Neural Network] ffnn_encoder_1--1x512-->stack_encoder((Stack)) ffnn_encoder_2--1x512-->stack_encoder ffnn_encoder_t--1x512-->stack_encoder ffnn_encoder_n--1x512-->stack_encoder end stack_encoder--nx512-->encoder_output[Encoder Output] end subgraph Decoder decoder_target[Decoder Target]--mx512-->pos_encoding_2[Positional Encoding] pos_encoding_2--mx512-->qkv_decoder_1[QKV Layer] qkv_decoder_1--Q=mx64-->multihead_attn_masked_1[Attention 1] qkv_decoder_1--K=mx64-->multihead_attn_masked_1 qkv_decoder_1--V=mx64-->multihead_attn_masked_1 qkv_decoder_1--Q=mx64-->multihead_attn_masked_2[Attention 2] qkv_decoder_1--K=mx64-->multihead_attn_masked_2 qkv_decoder_1--V=mx64-->multihead_attn_masked_2 qkv_decoder_1--Q=mx64-->multihead_attn_masked_3[Attention 3] qkv_decoder_1--K=mx64-->multihead_attn_masked_3 qkv_decoder_1--V=mx64-->multihead_attn_masked_3 qkv_decoder_1--Q=mx64-->multihead_attn_masked_4[Attention 4] qkv_decoder_1--K=mx64-->multihead_attn_masked_4 qkv_decoder_1--V=mx64-->multihead_attn_masked_4 qkv_decoder_1--Q=mx64-->multihead_attn_masked_5[Attention 5] qkv_decoder_1--K=mx64-->multihead_attn_masked_5 qkv_decoder_1--V=mx64-->multihead_attn_masked_5 qkv_decoder_1--Q=mx64-->multihead_attn_masked_6[Attention 6] qkv_decoder_1--K=mx64-->multihead_attn_masked_6 qkv_decoder_1--V=mx64-->multihead_attn_masked_6 qkv_decoder_1--Q=mx64-->multihead_attn_masked_7[Attention 7] qkv_decoder_1--K=mx64-->multihead_attn_masked_7 qkv_decoder_1--V=mx64-->multihead_attn_masked_7 qkv_decoder_1--Q=mx64-->multihead_attn_masked_8[Attention 8] qkv_decoder_1--K=mx64-->multihead_attn_masked_8 qkv_decoder_1--V=mx64-->multihead_attn_masked_8 subgraph DecoderMaskedMultiheadAttention[Decoder Masked Multihead Attention] multihead_attn_masked_1--mx64-->concat_masked((Concatenate)) multihead_attn_masked_2--mx64-->concat_masked multihead_attn_masked_3--mx64-->concat_masked multihead_attn_masked_4--mx64-->concat_masked multihead_attn_masked_5--mx64-->concat_masked multihead_attn_masked_6--mx64-->concat_masked multihead_attn_masked_7--mx64-->concat_masked multihead_attn_masked_8--mx64-->concat_masked end concat_masked--mx512-->linear_reproject_masked[Linear Reprojection] linear_reproject_masked--1x512-->ffnn_encoder_1_masked[FFNN 1] linear_reproject_masked--1x512-->ffnn_encoder_2_masked[FFNN 2] linear_reproject_masked--1x512-->ffnn_encoder_t_masked[FFNN x] linear_reproject_masked--1x512-->ffnn_encoder_n_masked[FFNN n] subgraph FfnnEncoderMasked[Feed Forward Neural Network] ffnn_encoder_1_masked--1x512-->stack_decoder_masked((Stack)) ffnn_encoder_2_masked--1x512-->stack_decoder_masked ffnn_encoder_t_masked--1x512-->stack_decoder_masked ffnn_encoder_n_masked--1x512-->stack_decoder_masked end stack_decoder_masked--mx512-->query_project[Query Projection] encoder_output--nx512-->kv_project_decoder[Key-Value Projection] query_project--Q=mx64-->multihead_attn_unmasked_1[Attention 1] kv_project_decoder--K=nx64-->multihead_attn_unmasked_1 kv_project_decoder--V=nx64-->multihead_attn_unmasked_1 query_project--Q=mx64-->multihead_attn_unmasked_2[Attention 2] kv_project_decoder--K=nx64-->multihead_attn_unmasked_2 kv_project_decoder--V=nx64-->multihead_attn_unmasked_2 query_project--Q=mx64-->multihead_attn_unmasked_3[Attention 3] kv_project_decoder--K=nx64-->multihead_attn_unmasked_3 kv_project_decoder--V=nx64-->multihead_attn_unmasked_3 query_project--Q=mx64-->multihead_attn_unmasked_4[Attention 4] kv_project_decoder--K=nx64-->multihead_attn_unmasked_4 kv_project_decoder--V=nx64-->multihead_attn_unmasked_4 query_project--Q=mx64-->multihead_attn_unmasked_5[Attention 5] kv_project_decoder--K=nx64-->multihead_attn_unmasked_5 kv_project_decoder--V=nx64-->multihead_attn_unmasked_5 query_project--Q=mx64-->multihead_attn_unmasked_6[Attention 6] kv_project_decoder--K=nx64-->multihead_attn_unmasked_6 kv_project_decoder--V=nx64-->multihead_attn_unmasked_6 query_project--Q=mx64-->multihead_attn_unmasked_7[Attention 7] kv_project_decoder--K=nx64-->multihead_attn_unmasked_7 kv_project_decoder--V=nx64-->multihead_attn_unmasked_7 query_project--Q=mx64-->multihead_attn_unmasked_8[Multihead Attention 8] kv_project_decoder--K=nx64-->multihead_attn_unmasked_8 kv_project_decoder--V=nx64-->multihead_attn_unmasked_8 subgraph DecoderUnmaskedMultiheadAttention[Decoder Unmasked Multihead Attention] multihead_attn_unmasked_1--mx64-->concat_unmasked((Concatenate)) multihead_attn_unmasked_2--mx64-->concat_unmasked multihead_attn_unmasked_3--mx64-->concat_unmasked multihead_attn_unmasked_4--mx64-->concat_unmasked multihead_attn_unmasked_5--mx64-->concat_unmasked multihead_attn_unmasked_6--mx64-->concat_unmasked multihead_attn_unmasked_7--mx64-->concat_unmasked multihead_attn_unmasked_8--mx64-->concat_unmasked end concat_unmasked--mx512-->linear_reproject_unmasked[Linear Reprojection] linear_reproject_unmasked--1x512-->ffnn_decoder_1_unmasked[FFNN 1] linear_reproject_unmasked--1x512-->ffnn_decoder_2_unmasked[FFNN 2] linear_reproject_unmasked--1x512-->ffnn_decoder_t_unmasked[FFNN x] linear_reproject_unmasked--1x512-->ffnn_decoder_n_unmasked[FFNN n] subgraph FfnnDecoderUnmasked[Feed Forward Neural Networks] ffnn_decoder_1_unmasked--1x512-->stack_decoder_unmasked((Stack)) ffnn_decoder_2_unmasked--1x512-->stack_decoder_unmasked ffnn_decoder_t_unmasked--1x512-->stack_decoder_unmasked ffnn_decoder_n_unmasked--1x512-->stack_decoder_unmasked end end stack_decoder_unmasked--mx512-->linear[Linear=512xV] subgraph OutputLayer[Output Layer] linear--mxV-->softmax[Softmax] softmax--mxV-->select_max_probabilities[Select Maximum Probability Token for each Position] select_max_probabilities--1xm-->transformer_output[Transformer Output] end

Notes on the Code

• The last word in the output is added to the output buffer, during inference.
• The encoder output is injected directly into the sublayer of every Decoder. To build up the chain of Decoders in PyTorch, so that we can put the full stack inside a Sequential block, we simply inject the encoder output to the root Decoder, and have it output the encoder output (together with the actual Decoder output) as part of the Decoder’s actual output to make it easy for the next Decoder in the stack to consume the Encoder and Decoder outputs.
• The code does not set up parameters in a form suitable for optimisation yet. There are several Module-subclasses which are really only there for the convenience of not having to call the forward() methods explicitly. In a sequel, we will collapse most of the parameters to be part of only a couple of Module subclasses.

• The class diagram is shown above. The composition hierarchy is quite straightforward, though there are some associations missing because of the shortcomings of the tool used to generate this (Pyreverse).
  • Specifically, DecoderStack contains a bunch of Decoders, and EncoderStack contains a bunch of Encoders.
  • qkv_source and masked_qkv_source contain instances of SingleSourceQKVLayer.
  • unmasked_qkv_source contains an instance of MultiSourceQKVLayer.
  • Some of the members in the classes are repeated because of Pyreverse duplicating the information from type hints.
• More notes can be found in the source itself.

Conclusion

We have built and tested the basic Transformer architecture. However, we still need to do the following:

• Build a proper vocabulary. Our current vocabulary is hard-coded, and contains random vectors.
• Several tensors are reused as parameters. Some of these need to be separate parameters.
• There are several Module subclasses. For optimisation, we will need to centralise where we register our parameters.
• We still need to train the Transformer.

All of the above, we will work on in the sequel to this post.

References

• Attention is All You Need
• Formal Algorithms for Transformers
• The Illustrated Transformer
• Transformers Explained Visually: Part 1
• Transformers Explained Visually: Part 2

---