Transformer Architecture: Blocks, Variants, and Costs

A Transformer architecture is a stack of neural-network blocks that updates token or patch representations with attention and position-wise feed-forward transformations. Residual paths preserve a shared representation stream, normalization controls activation scale, positional signals encode order, and masks define which positions may exchange information.

Attention is only one operation inside the architecture. A deployable Transformer must also specify its input representation, block order, normalization placement, output head, training objective, and inference state. Those choices determine what information can flow, what the model can predict, and how much memory and latency the system requires.

Key Takeaways

• A Transformer block preserves a hidden-state shape such as [batch, tokens, d_model] while attention mixes information across positions and the feed-forward network transforms each position independently.
• The original 2017 Transformer was an encoder-decoder model with post-layer normalization, sinusoidal positions, ReLU feed-forward layers, and multi-head attention. Those details are not universal requirements.
• Encoder-only, decoder-only, and encoder-decoder models differ primarily in visibility, objectives, and output flow, not in a simple "understanding versus generation" slogan.
• Training a causal model can process known target positions in parallel. Autoregressive decoding remains sequential because each next token depends on the token just selected.
• Dense attention is not the whole cost. Projections and feed-forward layers scale linearly with sequence length but quadratically with model width; KV-cache memory scales with tokens, layers, KV heads, head size, batch, and precision.
• Correctness tests should cover masks, positions, cached decoding, tokenizer/checkpoint compatibility, numerical finiteness, and task behavior.

What Makes a Transformer a Transformer

A Transformer combines an information-routing operation with a per-position transformation inside a residual stack. The original Attention Is All You Need model used an encoder and an autoregressive decoder for machine translation. Later models retained the block pattern while changing visibility, objectives, normalization, position handling, activations, and attention head layout.

For hidden states X with shape [B, T, D], a common pre-normalized block can be written as:

H = X + Attention(Norm(X), mask, position)
Y = H + FFN(Norm(H))

The original paper instead used post-normalization:

H = Norm(X + Attention(X, mask, position))
Y = Norm(H + FFN(H))

Both preserve the outer shape [B, T, D], which makes blocks stackable. They do not have identical optimization behavior. Xiong et al. analyzed why normalization placement changes gradient behavior at initialization; their results concern the studied architectures and training settings, not a guarantee that one layout wins for every model.

This page treats the attention mechanism as an already defined primitive. The separate attention guide covers QKV shapes, mask polarity, stable Softmax, MHA/MQA/GQA, and kernel behavior in detail.

Inside One Transformer Block

A useful implementation model is "read from the residual stream, compute a contribution, then add it back." Attention moves information between permitted positions. The feed-forward network changes each position's features with the same learned function.

Attention Sublayer

For one head, scaled dot-product attention is:

Attention(Q, K, V, M) = softmax((QK^T / sqrt(d_head)) + M) V

Multi-head attention learns separate projections, concatenates head outputs, and applies an output projection. A head has no guaranteed human-readable role. Claims such as "head 3 is syntax" require interventions or controlled analysis; an attention map alone is not a causal explanation.

The mask is part of the architecture:

• bidirectional self-attention exposes all permitted input positions;
• causal self-attention blocks future positions;
• cross-attention takes Queries from a target stream and Keys/Values from a source stream;
• padding and application masks remove invalid positions.

Feed-Forward Sublayer

The feed-forward network, or FFN, is applied independently to every position:

FFN(x) = activation(x W1 + b1) W2 + b2

The original model used ReLU. Later designs use alternatives such as GELU or gated variants. GLU Variants Improve Transformer reports improvements for several gated FFNs in its sequence-to-sequence experiments, but that result does not make a specific activation universally optimal.

Attention often receives most of the conceptual attention, yet the FFN can hold a large fraction of a dense model's parameters and matrix multiplication. Capacity and latency planning must include both sublayers.

Residual Connections and Normalization

Residual connections let each sublayer contribute an update without replacing the entire representation. Normalization placement affects optimization, checkpoint compatibility, and the exact forward function.

LayerNorm centers and rescales features. RMSNorm only rescales by root-mean-square magnitude; the RMSNorm paper evaluates its efficiency and task behavior under specific models. Treat LayerNorm versus RMSNorm, pre-norm versus post-norm, epsilon, and affine parameters as checkpoint contracts rather than interchangeable toggles.

Positional Information Is a Model Contract

Attention over Q/K/V alone does not encode sequence order. A Transformer therefore adds or applies positional information somewhere in the representation or attention calculation.

The original Transformer offered sinusoidal and learned absolute alternatives. BERT used learned token, segment, and position embeddings. T5 used a relative position bias. RoFormer introduced RoPE as rotations that encode absolute positions while making relative displacement visible in attention.

Changing a position scheme at runtime is not a cosmetic optimization. Position IDs, scaling rules, cached offsets, and maximum trained context must remain consistent. A model accepting a longer tensor does not prove that it uses the extra context reliably.

Encoder, Decoder, and Encoder-Decoder Variants

The three major Transformer families define different information flows. Architecture names are useful only when paired with the actual mask, objective, and output head.

BERT is evidence for a bidirectional encoder trained with masked-token prediction and a historical next-sentence objective. It is not proof that every encoder-only model uses BERT's objective. T5 demonstrates an encoder-decoder text-to-text framework and span corruption, not a universal recipe for every conditional task.

Decoder-only models can also classify, extract, or embed when their objective and output handling support those tasks. Encoder-only models can contribute to generation systems as retrievers or encoders. Select by the required information flow and measured system behavior, not by a rigid task label.

Training, Prefill, and Decode Are Different Workloads

Causal Transformer training and generation execute the same learned block under different data availability.

During teacher-forced training, the full target sequence is known. A causal mask prevents each position from reading future targets, but kernels can still process many positions in parallel. During inference:

1. Prefill processes the known prompt and creates layer-wise Keys and Values.
2. Decode consumes the selected token, reuses cached K/V state, and produces logits for the next token.
3. Sampling or decoding policy selects a token, creating a dependency before the next step.

A KV cache removes repeated K/V projection for prior tokens. It does not make autoregressive token selection parallel, and it introduces memory allocation, position, eviction, reuse, and tenant-isolation requirements. Hugging Face's cache documentation notes that cache position and attention-mask length must account for both past and current tokens and that caching is intended for inference.

Compute, Parameters, and KV-Cache Memory

Transformer cost depends on both sequence length and width. Saying only "Transformers are O(n squared)" hides important terms and can lead to poor capacity estimates.

For a dense block with sequence length T, model width D, and FFN width F:

The factors and fused operations vary by implementation. FlashAttention reduces intermediate memory traffic without changing dense attention's pairwise arithmetic. Sparse, local, recurrent, or state-space designs change the operation and must be evaluated under their own quality and hardware constraints.

A Runnable Go Architecture Audit

This dependency-free Go program validates a decoder configuration and estimates bias-free dense-block parameters plus KV-cache capacity. It models a SwiGLU-style FFN with three matrices and excludes embeddings, output heads, biases, Mixture-of-Experts layers, optimizer state, activations, and allocator overhead.

package main

import (
	"errors"
	"fmt"
)

type Config struct {
	Layers     int64
	ModelDim   int64
	FFNDim     int64
	QueryHeads int64
	KVHeads    int64
}

func (c Config) validate() error {
	if c.Layers <= 0 || c.ModelDim <= 0 || c.FFNDim <= 0 {
		return errors.New("layers and dimensions must be positive")
	}
	if c.QueryHeads <= 0 || c.KVHeads <= 0 {
		return errors.New("head counts must be positive")
	}
	if c.ModelDim%c.QueryHeads != 0 {
		return errors.New("model dimension must be divisible by query heads")
	}
	if c.QueryHeads%c.KVHeads != 0 {
		return errors.New("query heads must be divisible by KV heads")
	}
	return nil
}

func (c Config) headDim() int64 {
	return c.ModelDim / c.QueryHeads
}

func (c Config) blockParameters() int64 {
	kvWidth := c.KVHeads * c.headDim()
	attention := c.ModelDim*c.ModelDim +
		2*c.ModelDim*kvWidth +
		c.ModelDim*c.ModelDim
	swigluFFN := 3 * c.ModelDim * c.FFNDim
	twoRMSNorms := 2 * c.ModelDim
	return attention + swigluFFN + twoRMSNorms
}

func (c Config) kvCacheBytes(batch, tokens, bytesPerScalar int64) int64 {
	return 2 * batch * tokens * c.Layers *
		c.KVHeads * c.headDim() * bytesPerScalar
}

func gib(bytes int64) float64 {
	return float64(bytes) / (1 << 30)
}

func main() {
	config := Config{
		Layers:     32,
		ModelDim:   4096,
		FFNDim:     11008,
		QueryHeads: 32,
		KVHeads:    8,
	}
	if err := config.validate(); err != nil {
		panic(err)
	}

	fmt.Printf("head dimension: %d\n", config.headDim())
	fmt.Printf("parameters per block: %d\n", config.blockParameters())
	fmt.Printf(
		"KV cache: %.2f GiB\n",
		gib(config.kvCacheBytes(4, 8192, 2)),
	)
}

Expected output:

head dimension: 128
parameters per block: 177217536
KV cache: 4.00 GiB

The output is a planning estimate, not a profiler result. Real peak memory also includes weights, temporary workspaces, activations, fragmentation, runtime metadata, and potentially duplicated or quantized state.

Failure Modes That Architecture Diagrams Hide

A correct-looking block diagram can still produce invalid training or generation:

1. Future-token leakage: the causal mask is absent, shifted incorrectly, or has the wrong polarity.
2. Padding leakage: padding Keys remain visible or the loss includes padding targets.
3. Position drift: cached decoding restarts position IDs or applies RoPE with the wrong offset.
4. Checkpoint mismatch: tokenizer vocabulary, special tokens, normalization, head counts, or output tying differ from the weights.
5. Train/eval mismatch: dropout remains active, normalization state differs, or caching is enabled during training.
6. Cached-decode divergence: one-shot logits and token-by-token cached logits disagree beyond the expected numeric tolerance.
7. Fully masked rows or unstable logits: Softmax receives no valid Key or overflows in reduced precision.
8. Objective mismatch: the output head, label shift, and loss do not represent the intended task.

Framework APIs encode additional contracts. Current PyTorch TransformerEncoderLayer exposes batch_first and norm_first, warns that the module is a reference implementation with limited features relative to newer architectures, and warns that an incorrect causal hint can produce incorrect execution.

A Practical Validation Ladder

Validate a Transformer from local invariants to user-visible behavior:

1. Shape tests: assert every projection, residual addition, head split, and output dimension.
2. Mask tests: perturb a future or padded token and prove an earlier valid output does not change.
3. Numerical tests: reject NaN/Inf, compare stable reference attention, and define mixed-precision tolerances.
4. Cache parity: compare full-sequence next-token logits with iterative cached decoding at each position.
5. Tiny-batch test: deliberately overfit a tiny clean sample to expose wiring, loss, and label-shift errors.
6. Checkpoint contract: pin tokenizer, special-token IDs, architecture config, weight revision, dtype, and position settings.
7. Task evaluation: compare against a simple baseline with representative slices, not one aggregate score.
8. Serving evaluation: measure time to first token, inter-token latency, throughput, peak memory, context degradation, and failure behavior at realistic concurrency.

Architecture selection should follow this evidence. Compare encoder, decoder, encoder-decoder, recurrent, convolutional, state-space, retrieval, and hybrid systems under the same task and resource budget. The Mamba and state-space guide covers one alternative sequence-processing family; it does not imply that either architecture universally dominates.

Primary Sources

• Vaswani et al., Attention Is All You Need - original encoder-decoder Transformer, attention, FFN, residual, normalization, and sinusoidal position design.
• Devlin et al., BERT - bidirectional encoder and masked-token pretraining.
• Raffel et al., T5 - encoder-decoder text-to-text framework and relative position bias.
• Xiong et al., On Layer Normalization in the Transformer Architecture - analysis of pre-LN and post-LN gradient behavior.
• Su et al., RoFormer - rotary position embedding.
• Ainslie et al., Grouped-Query Attention - KV-head sharing and decoder-inference trade-offs.
• PyTorch TransformerEncoderLayer documentation - current shape and API contracts for the reference layer.

Conclusion

A Transformer is not just self-attention. It is a residual architecture whose masks, normalization, positional system, FFN, objective, and inference state jointly define behavior. Learn the original encoder-decoder design first, then treat modern decoder-only choices as explicit, checkpoint-specific variations.

For implementation work, keep shape and visibility contracts executable. For production work, separate prefill from decode, calculate KV-cache capacity before load testing, and require cached-decoding parity plus task-level evidence before trusting a model.