Transformer Data Flow

This closing note of the Transformers category traces one complete input all the way through the entire architecture — every component from this category, assembled into a single, continuous data flow, from raw tokens to a predicted next token.

The Complete Flow, Step by Step

1. Tokenization & embedding: raw input text is broken into tokens and converted into embedding vectors (covered fully in the NLP with Deep Learning category).
2. Positional encoding: position information (see Positional Encoding) is added to each token embedding.
3. Encoder stack: the combined embeddings pass through \(N\) encoder layers, each doing multi-head self-attention (with residual connection + layer norm) followed by a position-wise feed-forward network (with residual connection + layer norm), per Transformer Encoder.
4. Encoder output: one richly-contextualized vector per input token, ready to be queried by the decoder.
5. Decoder input: the target sequence generated so far (or, during training, the true target shifted right, connecting to Teacher Forcing) is embedded and positionally encoded, exactly like the encoder input.
6. Decoder stack: \(N\) decoder layers, each doing masked self-attention, then cross-attention over the encoder's output, then a feed-forward network — every sublayer wrapped in residual connections and layer normalization, per Transformer Decoder.
7. Output projection: the decoder's final output is projected (via one more linear layer) into a vector the size of the vocabulary, then passed through softmax to produce a probability distribution over the next token — exactly the categorical cross-entropy setup from Categorical Cross-Entropy.

Complete Diagram

Code — A Complete, Minimal Forward Pass

import torch
import torch.nn as nn

class MiniTransformer(nn.Module):
    def __init__(self, vocab_size, d_model=128, nhead=4, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = nn.Parameter(torch.randn(100, d_model))   # simplified, learned version
        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, batch_first=True)
        decoder_layer = nn.TransformerDecoderLayer(d_model, nhead, batch_first=True)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers)
        self.output_layer = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        src_emb = self.embedding(src) + self.pos_encoding[:src.size(1)]
        tgt_emb = self.embedding(tgt) + self.pos_encoding[:tgt.size(1)]

        encoder_output = self.encoder(src_emb)
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(1))
        decoder_output = self.decoder(tgt_emb, encoder_output, tgt_mask=tgt_mask)

        return self.output_layer(decoder_output)   # raw logits over the vocabulary

model = MiniTransformer(vocab_size=5000)
src = torch.randint(0, 5000, (1, 10))   # 10-token source sequence
tgt = torch.randint(0, 5000, (1, 6))     # 6 tokens generated/available so far
logits = model(src, tgt)
print(logits.shape)   # (1, 6, 5000) -- next-token probability logits at every decoder position

Common Mistakes

• Losing track of which parts of this flow run once (encoding the input) versus repeatedly (each autoregressive decoding step) — the encoder typically runs exactly once per input; the decoder conceptually runs once per generated token during inference, reusing the same fixed encoder output every time.
• Forgetting the final linear + softmax projection — the decoder's raw output is still in the \(d_{\text{model}}\)-dimensional space; it must be projected up to vocabulary size before it represents actual token probabilities.

Interview Relevance

Q: "Trace the complete data flow of a Transformer from input text to a predicted next token." A strong answer names every stage in order — tokenization/embedding, positional encoding, the encoder stack (self-attention + FFN, each with residual + norm), the decoder stack (masked self-attention, then cross-attention to the encoder output, then FFN, each with residual + norm), and finally the linear + softmax projection into a vocabulary-sized probability distribution — and can explain the role each stage plays, not just recite the names.

Key Takeaways — Transformers

• The Transformer replaces recurrence entirely with self-attention and feed-forward layers, unlocking full parallelization across the sequence dimension.
• Positional encoding restores the order-awareness self-attention lacks natively.
• Residual connections and layer normalization together make training very deep stacks of attention/feed-forward layers practically feasible.
• The decoder's masked self-attention prevents the model from "cheating" by seeing future tokens during training; cross-attention is the formal mechanism that resolves the original Seq2Seq context-vector bottleneck.
• Encoder-only, decoder-only, and full encoder-decoder variants adapt this same core toolkit to different tasks — understanding, generation, and sequence transduction respectively.

Next: NLP with Deep Learning covers how raw text actually becomes the token embeddings this category's Transformer consumes — tokenization, Word2Vec, GloVe, and the specific encoder-only and decoder-only architectures (BERT, GPT) built from everything in this category.

Practice Question

During inference, does the encoder run once per generated output token, or once total for the entire input sequence? Explain why.

Related DL Notes

• Transformer Encoder
• Transformer Decoder
• BERT