Transformer Encoder Stack

Suppose you want a single vector representation for the word "bank" that correctly reflects whether a sentence is about rivers or money. A Word Embeddings lookup table cannot do this. It gives "bank" the same vector regardless of context. You need a mechanism that lets every word's representation absorb information from every other word in the sentence. This must happen in both directions, simultaneously. That mechanism is the Transformer encoder stack: a tower of identical layers, each combining Multi-Head Attention with a position-wise feed-forward network. It turns a sequence of static token embeddings into a sequence of deeply contextualized ones. This assembly problem — how to stack sublayers so that a network trains stably at depths of 12, 24, or 96 layers — is not a minor engineering footnote. Before residual connections and layer normalization were combined with attention in the right order, deep sequence models suffered from vanishing gradients and unstable optimization. The specific recipe used in the Transformer encoder — attention sublayer, add & norm, feed-forward sublayer, add & norm, repeated identically — makes it possible to train networks deep enough. Such networks can capture long-range syntactic and semantic dependencies. At the same time, they remain numerically well-behaved. The encoder stack is the computational heart of models like BERT, RoBERTa, and countless sentence-embedding and retrieval systems. Unlike a decoder, which must not look at future tokens, the encoder is fully bidirectional. Every token attends to every other token, left and right, in a single pass. This bidirectionality is precisely what makes masked-language-model pretraining possible. It is why encoder stacks are the workhorse behind tasks like classification, tagging, and semantic search. They are not suited to open-ended text generation. This article builds the encoder stack from its sublayers upward. It covers how a single layer is wired, and why the residual-and-norm pattern is stacked to arbitrary depth without collapsing. It also covers the exact data flow through the stack, with a full pseudocode implementation. There is a worked numerical trace through two layers. Finally, it traces the historical path from the original 2017 paper to today's deep bidirectional encoders.

Key takeaways

The encoder stack is $N$ identical layers, each with a multi-head self-attention sublayer and a feed-forward sublayer, wrapped in residual connections and layer normalization.
Self-attention in the encoder is bidirectional — no causal mask — so every position sees the whole sequence at once.
Residual connections plus layer normalization are what allow stacking dozens of layers without gradient vanishing or exploding.
Positional encodings are injected only once, at the input, since attention and the FFN are otherwise permutation-equivariant.
Each layer costs $O(n^2 d + n d^2)$ time per sequence of length $n$ and model dimension $d$ ; the stack costs $N$ times that.
The design is the shared backbone of BERT-style pretraining, sentence encoders, and the encoder half of the full encoder-decoder Transformer.

The core insight

An encoder stack is $N$ copies of "bidirectional self-attention + feed-forward, each wrapped in residual+norm," repeated to build progressively richer, fully contextualized token representations.

The idea

A Transformer encoder stack is a sequence of $N$ identical encoder layers applied one after another to a sequence of input vectors. Each encoder layer contains exactly two sublayers: a multi-head self-attention sublayer, which lets every position mix information from every other position, and a position-wise feed-forward network, which transforms each position's vector independently and identically. Both sublayers are wrapped by the same structural pattern: compute the sublayer's output, add it back to its input (a residual connection), and normalize the result (layer normalization). This wrapping pattern — not the attention mechanism alone — is what makes the "stack" part of the name meaningful: it is the reason 12, 24, or even 48 such layers can be composed without the signal degrading. Formally, if $H^{(0)} \in \mathbb{R}^{n \times d}$ is the matrix of $n$ input token vectors (already summed with Positional Encodings since attention has no inherent notion of order), the stack computes

H^{(\ell)} = \mathrm{EncoderLayer}_\ell(H^{(\ell-1)}), \qquad \ell = 1, \dots, N,

and the final output $H^{(N)}$ is the sequence of contextualized representations handed to a downstream head — a classifier, a tagger, or (in the encoder-decoder setting) the decoder's cross-attention. Crucially, self-attention here is unmasked: position $i$ may attend to position $j$ for any $j$ , past or future. This is what distinguishes an encoder from a decoder, whose self-attention is causally masked so that position $i$ only attends to $j \le i$ . Bidirectionality is the encoder's defining property and the reason encoder-only models are trained with masked-language-modeling objectives rather than left-to-right next-token prediction. Each encoder layer is parameter-independent of the others — layer 3 has its own attention and FFN weights, distinct from layer 4's — but architecturally identical, which is what makes "stack $N$ of them" a well-defined, simple recipe rather than $N$ bespoke designs. The dimension $d$ (the model width) stays constant across the entire stack: every sublayer takes in and produces vectors of the same dimension $d$ , a design choice that is what makes residual addition ( $x + \mathrm{Sublayer}(x)$ ) well-typed at every layer.

How it works

Step 1 — Input preparation. Tokenize the input, map tokens to embeddings via an embedding matrix, and add positional encodings (sinusoidal or learned) so the model can distinguish token order:

H^{(0)}_i = \mathrm{Embed}(x_i) + \mathrm{PosEnc}(i).

Step 2 — Repeat $N$ times, per layer $\ell$ : 1. Self-attention sublayer. Compute multi-head self-attention over the entire current representation $H^{(\ell-1)}$ — queries, keys, and values all derive from the same tensor, and no positions are masked out. 2. Add & Norm. Add the sublayer's output back to its input (residual) and apply layer normalization. 3. Feed-forward sublayer. Apply a two-layer position-wise MLP (typically $d \to 4d \to d$ with a ReLU or GELU nonlinearity) independently to each position's vector. 4. Add & Norm. Again, residual-add and normalize. The layer's output becomes the next layer's input. After the final layer, the stack outputs $H^{(N)}$ , one contextualized vector per input token. Two placements of layer normalization exist in practice. Post-LN (the original 2017 design) normalizes after the residual addition: $\mathrm{LN}(x + \mathrm{Sublayer}(x))$ . Pre-LN (used in most modern large models) normalizes before the sublayer is applied: $x + \mathrm{Sublayer}(\mathrm{LN}(x))$ . Pre-LN keeps an unimpeded residual path directly from input to output at every layer, which empirically stabilizes training at large depth, at a mild cost in final representational sharpness; Post-LN can achieve marginally better final quality but is harder to train deep without learning-rate warm-up. ``` function ENCODER_STACK(tokens, N, d_model, num_heads, d_ff):

# tokens: sequence of length n (token ids) n = length(tokens) X = EmbeddingLookup(tokens) # n x d_model X = X + PositionalEncoding(n, d_model) # inject order information once X = Dropout(X) for layer_index in 1 .. N:

# ---- Sublayer 1: bidirectional multi-head self-attention ---- Y = LayerNorm(X) # Pre-LN variant A = MultiHeadAttention(query=Y, key=Y, value=Y, mask=NONE) X = X + Dropout(A) # residual connection

# ---- Sublayer 2: position-wise feed-forward ---- Y = LayerNorm(X) F = FeedForward(Y) # F = W2 activation(W1 Y + b1) + b2 X = X + Dropout(F) # residual connection X = LayerNorm(X) # final norm (Pre-LN needs this) return X # n x d_model, contextualized function MULTI_HEAD_ATTENTION(query, key, value, mask):

# query, key, value: n x d_model (identical tensor here: self-attention) heads = [] for h in 1 .. num_heads: Q = query @ W_Q[h] # n x d_k K = key @ W_K[h] # n x d_k V = value @ W_V[h] # n x d_v scores = (Q @ transpose(K)) / sqrt(d_k) # n x n if mask is not NONE: scores = scores + mask # encoder: mask is NONE (bidirectional) weights = Softmax(scores, axis=-1) # n x n, rows sum to 1 heads.append(weights @ V) # n x d_v concat = Concatenate(heads, axis=-1) # n x (num_heads * d_v) return concat @ W_O # n x d_model function FEED_FORWARD(X): return (Activation(X @ W1 + b1)) @ W2 + b2 # d_model -> d_ff -> d_model ``` Note the mask argument: in an encoder, `mask` is always `NONE` — this single line is the entire structural difference between encoder self-attention and the causally masked self-attention used in decoders.

Figure

One encoder layer: attention and feed-forward sublayers, each wrapped in residual + layer norm

Complexity & correctness

Per-layer time complexity. Self-attention computes an $n \times n$ score matrix via $Q K^\top$ , costing $O(n^2 d)$ time (with $d$ split across heads, the total work across all heads is still $O(n^2 d)$ since $\sum_h d_k^{(h)} = d$ ), plus the softmax and weighted sum over values, also $O(n^2 d)$ . The feed-forward sublayer applies a $d \to d_{ff} \to d$ MLP independently to each of the $n$ positions, costing $O(n \, d \, d_{ff})$ ; with the standard choice $d_{ff} = 4d$ , this is $O(n d^2)$ . So one layer costs

O(n^2 d + n d^2)

time. For a stack of $N$ layers, total time is $O(N(n^2 d + n d^2))$ . Space is dominated by the attention score matrices, $O(n^2)$ per layer (or $O(N n^2)$ if all are retained for backpropagation), plus $O(nd)$ for activations at each layer.

Why $n^2$ , and when it dominates. The quadratic term arises because every pair of positions is compared — there is no way to know which pairs matter without computing all $n^2$ attention scores. When $n \ll d$ (short sequences, wide models — typical of BERT-base with $n \approx 128$ – $512$ , $d = 768$ ), the $n d^2$ feed-forward term often dominates in practice. When $n$ grows large (long-document or genomic models), the $n^2 d$ term dominates, motivating the sparse and linear-attention variants that trade exactness for sub-quadratic scaling.

Correctness argument — why stacking preserves signal. The central risk in composing $N$ nonlinear transformations is that gradients either vanish (if each layer shrinks the signal) or explode (if each layer amplifies it), a phenomenon well documented in deep recurrent networks before residual architectures existed. The residual formulation $X_{\ell} = X_{\ell-1} + \mathrm{Sublayer}(X_{\ell-1})$ guarantees an identity path: the gradient of the loss with respect to $X_{\ell-1}$ contains a term $\partial \mathcal{L}/\partial X_\ell \cdot \mathbf{I}$ that passes through unchanged, regardless of how small the sublayer's own Jacobian becomes. By the chain rule across $N$ layers,

\frac{\partial X_N}{\partial X_0} = \prod_{\ell=1}^{N} \left( I + \frac{\partial \, \mathrm{Sublayer}_\ell}{\partial X_{\ell-1}} \right),

and expanding this product shows the identity term $I$ always survives as one additive component, so the gradient can never be driven to exactly zero purely by depth, unlike a plain product of Jacobians in a non-residual stack. Layer normalization complements this by rescaling each sublayer's input (or output, depending on Pre-/Post-LN) to have zero mean and unit variance across the feature dimension, which keeps activations in a numerically well-conditioned range at every depth and prevents the scale of representations from drifting layer to layer. Together, residual connections give a gradient guarantee and layer normalization gives an activation-scale guarantee — the combination is what licenses stacking to depths that a plain feed-forward composition of the same sublayers could not sustain.

Permutation equivariance and why position encodings are added once. Self-attention and the position-wise FFN are both permutation-equivariant: permuting the input rows permutes the output rows identically, with no other change (attention weights depend only on content, not index; the FFN is applied identically at each position). Consequently, a stack of such layers is itself permutation-equivariant, and the only place order information can enter is the initial addition of positional encodings before layer 1 — a fact that both explains why that addition is necessary and confirms that no later layer needs to reintroduce it.

Tradeoffs. The quadratic attention cost is the standard complaint against vanilla encoder stacks for very long sequences. Pre-LN improves trainability at depth but can slightly reduce peak achievable quality compared to a carefully warmed-up Post-LN model of the same size. Wider $d_{ff}$ increases the FFN's representational capacity but adds parameters and compute roughly linearly in $d_{ff}$ .

Worked example

Take a toy sequence of $n = 3$ tokens — "the", "cat", "sat" — with model dimension $d = 4$ (small enough to trace by hand) and a stack of $N = 2$ layers with a single attention head each.

Input. After embedding and positional-encoding addition, suppose:

token	$H^{(0)}$
the	$[1.0,\ 0.0,\ 1.0,\ 0.0]$
cat	$[0.0,\ 1.0,\ 0.0,\ 1.0]$
sat	$[1.0,\ 1.0,\ 0.0,\ 0.0]$

Layer 1, self-attention. With $W_Q = W_K = W_V = I$ (identity, for simplicity of hand-tracing), $Q = K = V = H^{(0)}$ . Compute raw scores $Q K^\top / \sqrt{d}$ with $\sqrt{4} = 2$ :

	the	cat	sat
the	$(1{+}0{+}1{+}0)/2 = 1.0$	$(0{+}0{+}0{+}0)/2 = 0.0$	$(1{+}0{+}0{+}0)/2 = 0.5$
cat	$0.0$	$(0{+}1{+}0{+}1)/2 = 1.0$	$(0{+}1{+}0{+}0)/2 = 0.5$
sat	$0.5$	$0.5$	$(1{+}1{+}0{+}0)/2 = 1.0$

Row-softmax gives, for the "the" row: $\mathrm{softmax}(1.0, 0.0, 0.5) \approx (0.44,\ 0.16,\ 0.40)$ . This says token "the" now attends 44% to itself, 16% to "cat", 40% to "sat" — notably, "the" does attend forward to "sat", something a causally masked decoder could never do. This is the concrete manifestation of bidirectionality. The attention output for "the" is the weighted sum $0.44 \cdot H^{(0)}_\text{the} + 0.16 \cdot H^{(0)}_\text{cat} + 0.40 \cdot H^{(0)}_\text{sat}$ $= 0.44[1,0,1,0] + 0.16[0,1,0,1] + 0.40[1,1,0,0]$ $= [0.84,\ 0.56,\ 0.44,\ 0.16]$ .

Add & norm. Residual: $[0.84,0.56,0.44,0.16] + [1.0,0.0,1.0,0.0] = [1.84,\ 0.56,\ 1.44,\ 0.16]$ . Layer-normalizing this vector (mean $\approx 0.75$ , subtract and rescale to unit variance) yields a vector with the same relative ordering but standardized scale — call it $\tilde X_\text{the}$ ; the exact numbers depend on the learned gain/bias, but the key qualitative point is that "the"'s representation has now absorbed a weighted mixture of "cat" and "sat," something the raw embedding never had.

Feed-forward. Each of the three (now attention-mixed, normalized) vectors passes independently through the same two-layer MLP, e.g. $d=4 \to d_{ff}=16 \to d=4$ with GELU — no mixing across tokens happens here, only a per-token nonlinear reshaping. After another residual-add and norm, this produces $H^{(1)}$ , the output of layer 1.

Layer 2. The entire process repeats with layer 2's own weight matrices ( $W_Q^{(2)}, W_K^{(2)}, \dots$ , independent of layer 1's), taking $H^{(1)}$ as input. Because "the"'s vector already encodes a mixture of "cat" and "sat" after layer 1, layer 2's attention scores are computed over these enriched vectors — meaning by the end of layer 2, "the"'s representation can reflect second-order relationships, e.g. information that flowed through "sat" from a fourth token, even though "the" never directly attended to that fourth token in layer 1. This is precisely how depth extends the effective receptive field of contextualization beyond what a single attention layer provides. The final $H^{(2)}$ is hand-off ready: for classification, a pooled or `[CLS]`-position vector is fed to a linear head; for tagging, every position's final vector is fed to a per-token classifier.

Common pitfalls & traps

Pitfall	Why it happens	Fix
Accidentally masking the encoder like a decoder	Copy-pasting decoder self-attention code, which applies a causal mask	Encoder self-attention must use `mask = NONE`; verify every position can attend to every other
Forgetting positional encodings entirely	Assuming attention "just knows" order	Add positional encodings once at the input; without them the stack is a bag-of-words model
Re-adding positional encodings at every layer	Confusing "inject once" with "inject per sublayer"	Add only before layer 1; layers preserve position information implicitly through their inputs
Mismatched residual dimensions	Changing $d_{model}$ mid-stack or between sublayer input/output	Keep every sublayer's output dimension exactly equal to $d_{model}$ so $x + \mathrm{Sublayer}(x)$ type-checks
Skipping layer normalization at very small learning rates and wondering why deep stacks diverge	Believing residual connections alone stabilize training	Pair residuals with normalization (Pre-LN or Post-LN) and appropriate warm-up; residuals alone do not bound activation scale
Applying dropout inconsistently between training and inference	Forgetting to disable dropout at inference time	Ensure dropout layers are gated by a training/eval flag, standard in any deep-learning framework

Pioneers & history

Origins · 1 minAshish Vaswani and how this was discovered

Ashish Vaswani

1980s–present · Google Brain research scientist, lead author of "Attention Is All You Need"

Vaswani and collaborators proposed replacing recurrence entirely with stacked self-attention, arguing that a purely attention-based architecture could match or exceed recurrent and convolutional sequence models while being dramatically more parallelizable — because every position in a layer can be processed simultaneously rather than sequentially, unlike an RNN's inherently sequential recurrence.

The encoder stack as described here is one half of the original Transformer, introduced alongside its causally masked decoder counterpart; encoder-only variants were soon separated out and became the dominant architecture for representation learning once paired with masked-language-model pretraining.

Origins · 1 minhow this was discovered

History

2015

— attention mechanisms popularized for sequence-to-sequence translation, layered on top of recurrent encoders.

2016

— layer normalization introduced as a normalization alternative to batch normalization, well suited to sequence models.

2017

— "Attention Is All You Need" introduces the full encoder-decoder Transformer, replacing recurrence with stacked self-attention and position-wise feed-forward sublayers.

2018

— BERT adopts the encoder stack alone, pretrained bidirectionally with a masked-language-modeling objective, establishing the encoder-only paradigm.

2020

s — Pre-LN variants, deeper stacks, and efficient/sparse attention approximations extend encoder stacks to much larger depths and longer sequences.

Where this matters

Encoder stacks are the backbone of BERT & Masked Language Modeling and its many descendants (RoBERTa, ALBERT, DeBERTa), which power search-engine query understanding, spam and toxicity classifiers, named-entity recognition, and semantic similarity systems used in retrieval-augmented pipelines. Sentence-embedding models (used for semantic search and clustering) run text through an encoder stack and pool the outputs into a single dense vector. In the full encoder-decoder Transformer used for machine translation and summarization, the encoder stack processes the source sequence once, and its final representations are read by every decoder layer via cross-attention. Vision Transformers repurpose the identical encoder-stack recipe — replacing tokens with image patches — showing that the architecture generalizes well beyond text.

Ask the tutor: How does swapping Post-LN for Pre-LN change the gradient flow analysis at layer 50 of a very deep encoder stack?

Why does the encoder stack use the same architecture at every layer instead of different designs per depth?

Using an architecturally identical block at every depth, with only the weights differing, keeps the design simple, easy to implement generically (`for layer in layers: x = layer(x)`), and easy to scale by simply increasing $N$ . Empirically, stacking identical residual blocks has proven sufficient: because each layer can learn to approximate an identity function when a transformation isn't needed (the residual path makes this trivial), the network can effectively "skip" layers it doesn't need rather than requiring hand-designed heterogeneity. This uniformity also makes weight-sharing variants (like ALBERT, which reuses one layer's weights across all depths) straightforward to define.

Does every encoder layer need both a residual connection and layer normalization, or would one alone suffice?

Both serve distinct, complementary purposes. The residual connection guarantees that gradients have an unimpeded identity path back to earlier layers, preventing vanishing gradients purely from depth. Layer normalization independently controls the scale of activations at each layer, preventing internal covariate-like drift that can destabilize training even when gradients aren't vanishing. Empirically, removing normalization from very deep stacks causes activations to grow or shrink uncontrollably across layers even with residuals present, while removing residuals from normalized stacks reintroduces the vanishing-gradient problem at depth. The combination, not either alone, is what licenses training at 12+ layers reliably.

How does the encoder stack differ from the decoder stack in the full Transformer?

Three differences matter. First, encoder self-attention is unmasked (bidirectional), while decoder self-attention is causally masked so token $i$ only attends to tokens $j \le i$ . Second, decoder layers include an additional cross-attention sublayer that attends over the encoder's final output, letting the decoder condition its generation on the source sequence. Third, encoders are typically used for representation tasks (classification, tagging, embedding), while decoders are used for autoregressive generation. See Full Transformer Architecture for how the two stacks are wired together in sequence-to-sequence models.

Why is the attention cost quadratic in sequence length, and does this make encoder stacks unusable for long documents?

Quadratic cost arises because computing attention scores requires comparing every pair of positions — an $n \times n$ score matrix, each entry costing $O(d)$ work, totals $O(n^2 d)$ . This is not a bug but a direct consequence of allowing unrestricted pairwise interaction between all tokens. For moderate sequence lengths (a few hundred to a couple thousand tokens), this cost is manageable given modern hardware. For much longer sequences, sparse-attention, low-rank, or chunked-attention variants trade some pairwise expressiveness for sub-quadratic scaling, but the standard dense encoder stack described here remains the default choice whenever sequence length is not extreme.

What role does the feed-forward sublayer play if attention already mixes information across tokens?

Self-attention is a linear combination of value vectors, followed by a linear output projection — it can reweight and blend information across positions, but by itself has limited capacity to apply nonlinear, position-specific transformations. The position-wise feed-forward sublayer supplies that nonlinear transformation, expanding each token's vector into a higher-dimensional space (commonly $4d$ ) and back down through a nonlinearity like GELU. Empirically, removing the FFN sublayer substantially degrades model quality even though attention remains — the two sublayers play genuinely complementary roles: attention routes information between positions, and the FFN reshapes information at each position.

Why must positional encodings, embeddings, and every sublayer output all share the same dimension d?

The entire stack is built around residual addition, $x + \mathrm{Sublayer}(x)$ , which is only well-defined when $x$ and $\mathrm{Sublayer}(x)$ have identical shape. Since the input to layer 1 already has dimension $d$ (embedding dimension), and every subsequent sublayer must preserve that dimension for the residual to type-check, $d$ is fixed as a single global hyperparameter for the entire stack. This is why multi-head attention's per-head dimension $d_k$ is chosen so that concatenating all heads and projecting back through $W_O$ restores exactly dimension $d$ , and why the feed-forward sublayer's output projection maps back from $d_{ff}$ to $d$ .

Is the encoder stack itself trained with a specific objective, or is it architecture-agnostic to the training task?

The encoder-stack architecture described here is purely a mechanism for producing contextualized representations; it does not itself define a training objective. It becomes useful only once paired with a task: masked-language-modeling pretraining (as in BERT & Masked Language Modeling), a supervised classification head, sequence-tagging losses, or contrastive objectives for sentence embeddings. This separation of architecture from objective is deliberate and is exactly what allows the same encoder-stack design to be reused across pretraining, fine-tuning, and entirely different downstream tasks without architectural change.