PyTorch Transformer Encoder: Architecture, Self-Attention, Encoder Blocks, and Implementation
A Transformer Encoder is a neural network architecture that converts an input sequence into contextual representations. It is one of the core building blocks of the original Transformer architecture and is widely used in models such as BERT, RoBERTa, and DeBERTa.
Unlike recurrent neural networks (RNNs), a Transformer encoder can process all positions in a sequence in parallel. Its self-attention mechanism allows each token to incorporate information from other relevant tokens in the sequence.
In this module, you will learn how a Transformer encoder works internally, how multi-head self-attention and feed-forward networks are combined, how residual connections and layer normalization improve training, and how to implement a Transformer encoder using PyTorch.
What You Will Learn
By the end of this module, you will understand:
- What a Transformer encoder is
- Transformer encoder architecture
- Encoder blocks and their components
- Self-attention
- Query, Key, and Value
- Multi-head self-attention
- Feed-forward networks
- Residual connections
- Layer normalization
- Attention masks and padding masks
- Positional information
- How to build an encoder block in PyTorch
- How to stack multiple encoder layers
- How to use
nn.TransformerEncoder - How token embeddings become contextual representations
- Best practices for implementing Transformer encoders
What Is a Transformer Encoder?
A Transformer encoder processes an input sequence and produces a contextual representation for every input position.
For example:
1The cat sat on the mat
The representation of the word sat can incorporate information from other words in the sentence.
The encoder therefore does not treat each token independently. Instead, self-attention allows the representation at each position to be influenced by other positions.
A simplified workflow is:
1Input Token IDs 2 │ 3 ▼ 4Token Embeddings 5 │ 6 ▼ 7Positional Information 8 │ 9 ▼ 10Transformer Encoder Layers 11 │ 12 ▼ 13Contextual Representations
Why Use a Transformer Encoder?
Traditional recurrent architectures process tokens sequentially:
1Token 1 2 ↓ 3Token 2 4 ↓ 5Token 3 6 ↓ 7Token 4
This can make parallel computation difficult.
A Transformer encoder instead processes the sequence using attention:
1Token 1 ─┐ 2Token 2 ─┤ 3Token 3 ─┼──► Self-Attention 4Token 4 ─┤ 5Token 5 ─┘
The attention mechanism allows information to flow between positions while making the computation highly parallelizable.
Transformer Encoder Architecture
A typical encoder can be represented as:
1 Input Tokens 2 │ 3 ▼ 4 Token Embedding 5 │ 6 ▼ 7 Positional Information 8 │ 9 ▼ 10 ┌──────────────────────────┐ 11 │ Encoder Layer 1 │ 12 └──────────────────────────┘ 13 │ 14 ▼ 15 ┌──────────────────────────┐ 16 │ Encoder Layer 2 │ 17 └──────────────────────────┘ 18 │ 19 ▼ 20 ... 21 │ 22 ▼ 23 ┌──────────────────────────┐ 24 │ Encoder Layer N │ 25 └──────────────────────────┘ 26 │ 27 ▼ 28 Contextual Features
Each encoder layer transforms the representations produced by the previous layer.
Transformer Encoder Block
A standard Transformer encoder layer contains two main sub-layers:
- Multi-head self-attention
- Position-wise feed-forward network
Residual connections and layer normalization are applied around these sub-layers.
A common post-normalization representation is:
1Input 2 │ 3 ▼ 4Multi-Head Self-Attention 5 │ 6 ▼ 7Dropout 8 │ 9 ├───────────────┐ 10 │ │ 11 ▼ │ 12 Add ◄────────────┘ 13 │ 14 ▼ 15LayerNorm 16 │ 17 ▼ 18Feed-Forward Network 19 │ 20 ▼ 21Dropout 22 │ 23 ├───────────────┐ 24 │ │ 25 ▼ │ 26 Add ◄────────────┘ 27 │ 28 ▼ 29LayerNorm 30 │ 31 ▼ 32Output
Modern Transformer implementations may instead use pre-layer normalization, where normalization occurs before the attention and feed-forward sub-layers.
Self-Attention
Self-attention allows every token to interact with other tokens in the same sequence.
Consider:
1The cat sat on the mat
When processing sat, the model can assign attention to tokens such as:
1The 2cat 3sat 4on 5the 6mat
The attention weights determine how strongly information from each position contributes to the representation of the current position.
This allows the model to capture relationships that may be far apart in the sequence.
Query, Key, and Value
Self-attention uses three representations:
1Query (Q) 2Key (K) 3Value (V)
They are produced from the input representation using learned projections.
Conceptually:
1Input 2 │ 3 ├──► Query 4 │ 5 ├──► Key 6 │ 7 └──► Value
The attention mechanism compares queries with keys to determine attention scores and then uses those scores to combine the values.
A simplified scaled dot-product attention equation is:
1Attention(Q, K, V) 2= 3softmax(QKᵀ / √dₖ)V
Where:
Qrepresents queriesKrepresents keysVrepresents valuesdₖis the key dimension
Self-Attention Workflow
The complete process can be summarized as:
1Input Representations 2 │ 3 ├──────────► Q 4 ├──────────► K 5 └──────────► V 6 │ 7 ▼ 8 Q × Kᵀ 9 │ 10 ▼ 11 Scale by √dₖ 12 │ 13 ▼ 14 Softmax 15 │ 16 ▼ 17 Attention Weights 18 │ 19 ▼ 20 × Values 21 │ 22 ▼ 23 Attention Output
Multi-Head Self-Attention
Instead of calculating attention using a single representation, Transformers divide the representation into multiple attention heads.
For example:
1 Input 2 │ 3 ┌─────────┼─────────┐ 4 ▼ ▼ ▼ 5 Head 1 Head 2 Head 3 ... 6 │ │ │ 7 └─────────┼─────────┘ 8 ▼ 9 Concatenate 10 │ 11 ▼ 12 Linear Projection 13 │ 14 ▼ 15 Output
Different heads can learn different relationships between tokens.
For example, one head may learn relationships related to syntax while another may focus on semantic or positional relationships. The exact behavior of individual attention heads is learned during training rather than explicitly assigned.
Multi-Head Attention in PyTorch
PyTorch provides nn.MultiheadAttention.
1import torch 2import torch.nn as nn 3 4attention = nn.MultiheadAttention( 5 embed_dim=512, 6 num_heads=8, 7 batch_first=True 8) 9 10x = torch.randn( 11 2, 12 12, 13 512 14) 15 16output, weights = attention( 17 x, 18 x, 19 x 20) 21 22print(output.shape) 23print(weights.shape)
Output:
1torch.Size([2, 12, 512]) 2torch.Size([2, 12, 12])
The dimensions represent:
12 → batch size 212 → sequence length 3512 → embedding dimension
The attention weights describe relationships between sequence positions, although their returned shape can depend on configuration such as whether attention weights are averaged across heads.
Attention Masks
Encoder self-attention is generally bidirectional, meaning a token can attend to other non-masked positions on both sides.
However, encoders commonly use a padding mask to prevent attention from being assigned to padding tokens.
For example:
1The cat sat <PAD> <PAD>
The <PAD> positions should not contribute meaningful information.
A padding mask can be represented as:
1padding_mask = tokens == pad_token_id
The exact mask argument depends on the PyTorch API being used.
Encoder Attention vs Decoder Causal Attention
It is important to distinguish encoder attention from decoder attention.
Encoder
The encoder normally allows bidirectional attention:
1Token 1 ←→ Token 2 ←→ Token 3 ←→ Token 4
A token can use information from both earlier and later positions, except where masking prevents it.
Autoregressive Decoder
A causal decoder restricts future positions:
1Token 1 2 ↓ 3Token 1 → Token 2 4 ↓ 5Token 1 → Token 2 → Token 3
Therefore, a standard encoder does not normally require the causal look-ahead mask used by an autoregressive decoder.
Feed-Forward Network
After self-attention, the encoder applies a position-wise feed-forward network.
A simplified architecture is:
1Input 2 │ 3 ▼ 4Linear 5 │ 6 ▼ 7Activation 8 │ 9 ▼ 10Dropout 11 │ 12 ▼ 13Linear 14 │ 15 ▼ 16Output
For example:
1512 2 ↓ 32048 4 ↓ 5512
The feed-forward network is applied independently to each sequence position while using the same learned parameters at every position.
Feed-Forward Network in PyTorch
1import torch 2import torch.nn as nn 3 4ffn = nn.Sequential( 5 nn.Linear( 6 512, 7 2048 8 ), 9 10 nn.GELU(), 11 12 nn.Dropout(0.1), 13 14 nn.Linear( 15 2048, 16 512 17 ) 18) 19 20x = torch.randn( 21 4, 22 16, 23 512 24) 25 26output = ffn(x) 27 28print(output.shape)
Output:
1torch.Size([4, 16, 512])
The sequence length and model dimension remain unchanged.
Layer Normalization
Layer normalization normalizes the feature dimensions of each token representation.
For a tensor shaped:
1(batch, sequence_length, d_model)
LayerNorm operates over the feature dimension specified by normalized_shape.
For example:
1layer_norm = nn.LayerNorm(512)
Example:
1import torch 2import torch.nn as nn 3 4layer_norm = nn.LayerNorm(512) 5 6x = torch.randn( 7 2, 8 10, 9 512 10) 11 12output = layer_norm(x) 13 14print(output.shape)
Output:
1torch.Size([2, 10, 512])
Layer normalization does not normalize across the batch in the same way as BatchNorm.
Why Layer Normalization Is Important
Layer normalization helps stabilize the values flowing through Transformer layers.
It is particularly useful in deep networks where repeated transformations can otherwise make optimization more difficult.
Important benefits include:
- More stable optimization
- Improved training behavior
- Better gradient propagation
- Independence from batch-level statistics
Residual Connections
Residual connections add the input of a sub-layer to its transformed output.
Conceptually:
1 ┌──────────────┐ 2 │ │ 3 ▼ │ 4Input ──► Sublayer ──► Add ──┘ 5 │ 6 ▼ 7 Output
Mathematically:
1Output = x + Sublayer(x)
Residual connections help preserve information and improve gradient flow through deep Transformer networks.
Residual Connection in PyTorch
1import torch 2import torch.nn as nn 3 4layer = nn.Linear( 5 512, 6 512 7) 8 9x = torch.randn( 10 2, 11 8, 12 512 13) 14 15output = x + layer(x) 16 17print(output.shape)
Output:
1torch.Size([2, 8, 512])
In a real Transformer encoder, the residual connection is combined with dropout and normalization.
Positional Information
Self-attention itself does not inherently provide the model with the order of tokens.
For example:
1The cat sat
and:
1cat The sat
contain the same token set but have different meanings.
Therefore, Transformer architectures need positional information so the model can distinguish token positions.
The original Transformer architecture used sinusoidal positional encodings.
Modern architectures may use other approaches, including:
- Learned positional embeddings
- Rotary Position Embeddings (RoPE)
- Relative position representations
- Other position-aware mechanisms
The positional mechanism depends on the model architecture.
Token Embeddings
Before entering the encoder, token IDs are converted into dense vectors.
For example:
1embedding = nn.Embedding( 2 10000, 3 512 4)
If the input is:
1(batch=2, sequence_length=30)
the embedding output becomes:
1(batch=2, sequence_length=30, d_model=512)
Example:
1tokens = torch.randint( 2 0, 3 10000, 4 (2, 30) 5) 6 7x = embedding(tokens) 8 9print(x.shape)
Output:
1torch.Size([2, 30, 512])
Building a Transformer Encoder Block
The following implementation demonstrates the main components of an encoder block.
1import torch 2import torch.nn as nn 3 4 5class EncoderBlock(nn.Module): 6 7 def __init__( 8 self, 9 d_model=512, 10 n_heads=8, 11 ff_dim=2048, 12 dropout=0.1 13 ): 14 super().__init__() 15 16 self.attention = nn.MultiheadAttention( 17 embed_dim=d_model, 18 num_heads=n_heads, 19 dropout=dropout, 20 batch_first=True 21 ) 22 23 self.norm1 = nn.LayerNorm( 24 d_model 25 ) 26 27 self.norm2 = nn.LayerNorm( 28 d_model 29 ) 30 31 self.ffn = nn.Sequential( 32 nn.Linear( 33 d_model, 34 ff_dim 35 ), 36 37 nn.GELU(), 38 39 nn.Dropout(dropout), 40 41 nn.Linear( 42 ff_dim, 43 d_model 44 ) 45 ) 46 47 self.dropout = nn.Dropout( 48 dropout 49 ) 50 51 def forward( 52 self, 53 x, 54 key_padding_mask=None 55 ): 56 57 attention_output, _ = self.attention( 58 x, 59 x, 60 x, 61 key_padding_mask=key_padding_mask 62 ) 63 64 x = self.norm1( 65 x + self.dropout( 66 attention_output 67 ) 68 ) 69 70 ffn_output = self.ffn(x) 71 72 x = self.norm2( 73 x + self.dropout( 74 ffn_output 75 ) 76 ) 77 78 return x
This example uses a post-norm arrangement similar to the original Transformer formulation.
Testing the Encoder Block
1encoder = EncoderBlock() 2 3x = torch.randn( 4 4, 5 15, 6 512 7) 8 9output = encoder(x) 10 11print(output.shape)
Output:
1torch.Size([4, 15, 512])
The encoder preserves:
1batch size 2sequence length 3model dimension
while transforming the representation at every layer.
Using a Padding Mask
Suppose the batch contains padded sequences.
1tokens = torch.tensor([ 2 [1, 2, 3, 4, 0, 0], 3 [5, 6, 7, 0, 0, 0] 4]) 5 6pad_token_id = 0 7 8padding_mask = ( 9 tokens == pad_token_id 10)
The result identifies positions that should be ignored as padding.
For nn.MultiheadAttention, this can be supplied through key_padding_mask.
1output, weights = attention( 2 x, 3 x, 4 x, 5 key_padding_mask=padding_mask 6)
The mask must have the appropriate shape for the batch and sequence dimensions.
Building Multiple Encoder Layers
A complete Transformer encoder usually stacks several encoder blocks.
1class TransformerEncoder(nn.Module): 2 3 def __init__( 4 self, 5 num_layers=6, 6 d_model=512, 7 n_heads=8, 8 ff_dim=2048, 9 dropout=0.1 10 ): 11 super().__init__() 12 13 self.layers = nn.ModuleList([ 14 EncoderBlock( 15 d_model=d_model, 16 n_heads=n_heads, 17 ff_dim=ff_dim, 18 dropout=dropout 19 ) 20 for _ in range(num_layers) 21 ]) 22 23 def forward( 24 self, 25 x, 26 key_padding_mask=None 27 ): 28 29 for layer in self.layers: 30 31 x = layer( 32 x, 33 key_padding_mask=key_padding_mask 34 ) 35 36 return x
The representation flows through each encoder layer:
1Input 2 │ 3 ▼ 4Encoder Layer 1 5 │ 6 ▼ 7Encoder Layer 2 8 │ 9 ▼ 10Encoder Layer 3 11 │ 12 ▼ 13... 14 │ 15 ▼ 16Encoder Layer N 17 │ 18 ▼ 19Contextual Representation
Testing the Encoder Stack
1encoder = TransformerEncoder( 2 num_layers=6 3) 4 5x = torch.randn( 6 2, 7 20, 8 512 9) 10 11output = encoder(x) 12 13print(output.shape)
Output:
1torch.Size([2, 20, 512])
Using PyTorch TransformerEncoderLayer
PyTorch provides a built-in encoder layer:
1import torch 2import torch.nn as nn 3 4encoder_layer = nn.TransformerEncoderLayer( 5 d_model=512, 6 nhead=8, 7 dim_feedforward=2048, 8 dropout=0.1, 9 batch_first=True 10) 11 12x = torch.randn( 13 2, 14 12, 15 512 16) 17 18output = encoder_layer(x) 19 20print(output.shape)
Output:
1torch.Size([2, 12, 512])
The built-in implementation is useful when you want a standard Transformer encoder layer without implementing the individual components yourself.
Building a Complete PyTorch Transformer Encoder
Multiple encoder layers can be combined using nn.TransformerEncoder.
1import torch 2import torch.nn as nn 3 4encoder_layer = nn.TransformerEncoderLayer( 5 d_model=512, 6 nhead=8, 7 dim_feedforward=2048, 8 dropout=0.1, 9 batch_first=True 10) 11 12encoder = nn.TransformerEncoder( 13 encoder_layer, 14 num_layers=6 15) 16 17x = torch.randn( 18 4, 19 25, 20 512 21) 22 23output = encoder(x) 24 25print(output.shape)
Output:
1torch.Size([4, 25, 512])
Building an Encoder With Token Embeddings
A more realistic model starts with token IDs and converts them into embeddings.
1import torch 2import torch.nn as nn 3 4 5class TransformerModel(nn.Module): 6 7 def __init__( 8 self, 9 vocab_size, 10 d_model=512, 11 n_heads=8, 12 ff_dim=2048, 13 num_layers=6, 14 dropout=0.1 15 ): 16 super().__init__() 17 18 self.embedding = nn.Embedding( 19 vocab_size, 20 d_model 21 ) 22 23 encoder_layer = nn.TransformerEncoderLayer( 24 d_model=d_model, 25 nhead=n_heads, 26 dim_feedforward=ff_dim, 27 dropout=dropout, 28 batch_first=True 29 ) 30 31 self.encoder = nn.TransformerEncoder( 32 encoder_layer, 33 num_layers=num_layers 34 ) 35 36 def forward( 37 self, 38 x, 39 padding_mask=None 40 ): 41 42 x = self.embedding(x) 43 44 return self.encoder( 45 x, 46 src_key_padding_mask=padding_mask 47 )
This example focuses on the encoder itself. A production model would also need an appropriate positional-information mechanism.
Testing the Transformer Model
1model = TransformerModel( 2 vocab_size=10000, 3 d_model=512 4) 5 6tokens = torch.randint( 7 0, 8 10000, 9 (2, 30) 10) 11 12output = model(tokens) 13 14print(output.shape)
Output:
1torch.Size([2, 30, 512])
The result contains one contextual representation for every input position.
Understanding the Encoder Output
Suppose:
1output.shape 2= 3(batch, sequence_length, d_model)
For:
1(2, 30, 512)
we have:
12 2→ Number of examples 3 430 5→ Number of tokens per sequence 6 7512 8→ Representation size for each token
For example:
1token_representation = output[0, 0] 2 3print(token_representation.shape)
Output:
1torch.Size([512])
This vector represents the contextual representation of the first token for the first example.
Encoder Workflow
The complete workflow can be summarized as:
1Input Token IDs 2 │ 3 ▼ 4Token Embedding 5 │ 6 ▼ 7Positional Information 8 │ 9 ▼ 10Multi-Head Self-Attention 11 │ 12 ▼ 13Residual + Normalization 14 │ 15 ▼ 16Feed-Forward Network 17 │ 18 ▼ 19Residual + Normalization 20 │ 21 ▼ 22Repeat N Layers 23 │ 24 ▼ 25Final Contextual Representations
Transformer Encoder Applications
Transformer encoders are especially useful when the model needs a rich representation of the complete input sequence.
Common applications include:
- Text classification
- Sentiment analysis
- Named Entity Recognition
- Token classification
- Sentence embeddings
- Semantic search
- Information retrieval
- Document understanding
- Question answering
- Feature extraction
Encoder-based architectures include:
- BERT
- RoBERTa
- DeBERTa
These models differ in training objectives and architectural details, but they are based on encoder-style Transformer representations.
Encoder vs Decoder
| Feature | Transformer Encoder | Autoregressive Transformer Decoder |
|---|---|---|
| Self-attention | Yes | Yes |
| Attention direction | Usually bidirectional | Causal |
| Future-token masking | Usually no | Yes |
| Padding mask | Commonly used | Commonly used |
| Cross-attention | No | Yes in encoder-decoder architectures |
| Main purpose | Input representation | Sequence generation |
| Example architectures | BERT, RoBERTa, DeBERTa | GPT, LLaMA |
The terminology is important: GPT and LLaMA are decoder-only language models, while T5 and BART use encoder-decoder architectures.
Common Transformer Encoder Mistakes
Forgetting Positional Information
Self-attention does not automatically encode token order, so the architecture needs a suitable positional mechanism.
Using a Causal Mask Unnecessarily
A standard encoder is normally bidirectional. Applying a causal mask changes its attention pattern and is generally not appropriate for a standard encoder architecture.
Ignoring Padding Tokens
Padded positions should normally be excluded from attention when processing variable-length batches.
Using an Invalid Number of Attention Heads
The model dimension must be compatible with the number of attention heads.
For example:
1d_model = 512 2n_heads = 8
gives:
1512 / 8 = 64
features per head.
Confusing Encoder Output With Class Predictions
The encoder normally produces contextual representations.
A task-specific head is then added for classification, token prediction, embeddings, or another downstream task.
For example:
1Encoder 2 │ 3 ▼ 4Contextual Representations 5 │ 6 ▼ 7Task-Specific Head 8 │ 9 ▼ 10Prediction
Best Practices
- Keep the embedding dimension compatible with the number of attention heads.
- Use an appropriate positional-information mechanism.
- Use padding masks when processing padded sequences.
- Do not use causal masking unless the architecture specifically requires it.
- Use residual connections around attention and feed-forward sub-layers.
- Use LayerNorm consistently with the chosen Transformer architecture.
- Use dropout where appropriate for regularization.
- Start with smaller models when learning or debugging Transformer implementations.
- Verify tensor shapes at every major stage.
- Use PyTorch's built-in Transformer modules when a custom implementation is not necessary.
- For very deep Transformers, understand the difference between pre-norm and post-norm architectures.
- Separate the encoder from task-specific prediction heads so the representation can be reused for different downstream tasks.
Practice Project: Build a Transformer Encoder
Create a small Transformer encoder using:
1Vocabulary Size = 5,000 2Model Dimension = 256 3Attention Heads = 8 4Encoder Layers = 4 5Sequence Length = 20 6Batch Size = 8
Step 1: Create Input Tokens
1tokens = torch.randint( 2 0, 3 5000, 4 (8, 20) 5)
Step 2: Build the Encoder
1model = TransformerModel( 2 vocab_size=5000, 3 d_model=256, 4 n_heads=8, 5 num_layers=4 6)
Step 3: Run the Forward Pass
1output = model(tokens) 2 3print(output.shape)
Expected output:
1torch.Size([8, 20, 256])
Step 4: Inspect One Token Representation
1print(output[0, 0])
The result is a vector containing 256 learned contextual features for that position.
Practice Exercises
Exercise 1: Change the Number of Layers
Try:
1num_layers=2
and then:
1num_layers=8
Compare the output shape and model parameter count.
Exercise 2: Change the Model Dimension
Try:
1d_model=256
and:
1d_model=512
Observe how the representation size changes.
Exercise 3: Add Padding
Create sequences with different effective lengths and use a padding mask.
Verify that padded positions are excluded from attention.
Exercise 4: Build a Classification Head
Take the encoder output and add a classification layer:
1Encoder 2 │ 3 ▼ 4Pooling / Selected Representation 5 │ 6 ▼ 7Linear Layer 8 │ 9 ▼ 10Class Logits
This is a useful next step for understanding how encoder models such as BERT are adapted to downstream tasks.
Module Summary
A Transformer encoder converts an input sequence into contextual representations using repeated self-attention and feed-forward transformations.
The core architecture is:
1Input 2 │ 3 ▼ 4Token Embeddings 5 │ 6 ▼ 7Positional Information 8 │ 9 ▼ 10Multi-Head Self-Attention 11 │ 12 ▼ 13Residual + LayerNorm 14 │ 15 ▼ 16Feed-Forward Network 17 │ 18 ▼ 19Residual + LayerNorm 20 │ 21 ▼ 22Repeated Encoder Layers 23 │ 24 ▼ 25Contextual Representations
The most important concepts to remember are:
- Self-attention allows tokens to exchange information with other positions.
- Multi-head attention performs attention through multiple learned projections.
- Feed-forward networks transform each token representation independently.
- Residual connections help information and gradients flow through deep networks.
- Layer normalization helps stabilize Transformer training.
- Positional information provides information about token order.
- Padding masks prevent padded positions from participating in attention.
- Encoder self-attention is normally bidirectional, unlike the causal attention used by autoregressive decoders.
Once you understand the Transformer encoder, the next step is to study the Transformer Decoder, including causal masking, cross-attention, autoregressive generation, teacher forcing, and the differences between decoder-only and encoder-decoder Transformer architectures.