Neural network architectures are built from components that transform representations and control how information flows through the model. This article collects short explanations of these components, starting with dropout.

Dropout

Dropout randomly sets part of an activation tensor to zero during training. It is a regularizer. The model cannot rely on every feature always being present, so it is pushed to learn less brittle representations.

During training, each value is kept with probability $1-p$ and dropped with probability $p$. The kept values are divided by $1-p$:

where $m$ is a random mask with values $0$ or $1$.

The scaling keeps the expected value the same. For one activation $x$:

That is why in evaluation dropout does nothing but in training samples a mask and rescales the kept activations.

In distributed training, replicated activations that must remain identical need the same dropout mask on every replica. Different shards of one activation instead use independent mask entries, so together they form one independent mask for the full tensor.

Dropout should usually come after the activation because some activations turn an input of zero into a nonzero output. For example, $\operatorname{sigmoid}(0)=0.5$, so masking a value before sigmoid does not actually remove it. Masking after the activation always makes the dropped feature zero. With ReLU the two placements can be equivalent because $\operatorname{ReLU}(0)=0$.

Applying dropout after the final logits usually does not make sense. Logits are the direct inputs to the loss. If we randomly zero them, we are corrupting the model's answer right before the objective reads it.

• Analysis of Dropout
• Improving neural networks by preventing co-adaptation of feature detectors, the original Dropout paper.

Residual connection

A residual block is a group of layers with a shortcut connection that adds the original input to the layers' output. If the layers compute $F(x)$, the block returns $F(x)+x$.

During the backward pass, the gradient can flow through both paths. Let $y=F(x)+x$ and let $J_F(x)$ be the Jacobian of $F$. Then

The first term is the gradient through $F$, and the second is the gradient through the identity shortcut.

Suppose the desired output for input $x$ is the underlying mapping $H(x)$. Instead of asking the stacked layers to learn $H(x)$ directly, a residual block lets them learn a residual function $F(x)$ relative to the input. The original mapping is then obtained by adding the input back:

The idea is that optimizing the residual mapping can be easier than optimizing the original, unreferenced mapping. If the identity mapping is already a good solution, so that $H(x)=x$, the stacked layers only need to push $F(x)$ toward zero. This can be easier than learning the identity function from scratch through multiple nonlinear layers. The Deep Residual Learning paper introduced this approach.

Attention

Attention produces each token's new representation as a weighted sum of the hidden states of other tokens. For token $i$, it has the form:

Here $v_j$ is the value representation of token $j$, and $\alpha_{ij}$ is the weight that token $i$ assigns to it. The weights are computed afresh for each input from query-key similarity scores and normalized with softmax. Training learns the query, key, and value projections that produce those scores, so the model learns how to assign useful weights rather than learning one fixed set of weights for every sequence.

Feed-forward network (FFN)

The feed-forward network is the Transformer block's per-token feature extractor. Attention mixes information across tokens, while the FFN transforms each token representation independently. A standard FFN expands the hidden size from $d$ to a larger intermediate size $d_{\text{ff}}$, applies a nonlinearity, then projects back to $d$:

In a Transformer block, the FFN follows multi-head attention. Attention gathers information from other tokens, and its output projection combines the attention heads. The FFN then mixes those combined features within each token and creates new nonlinear features for the next block.

The intermediate dimension is usually larger than the model dimension, often $4d$. Reducing it below $d$ would create a bottleneck that compresses the representation before the model can form many useful features. The up projection gives the nonlinearity more feature directions to combine, and the down projection returns the result to the fixed hidden size required by the residual stream.

The choice of roughly $4d$ is an empirical capacity–cost tradeoff. With $d$ input features, even considering pairwise relationships gives on the order of $d^2$ possible interactions. A network cannot explicitly represent every possible interaction, but a very narrow intermediate layer would leave it with too few nonlinear features to represent many useful combinations. Expanding the width gives the FFN room to develop more patterns before projecting them back, while keeping the parameter count and compute practical.

Multi-head attention (MHA)

MHA gives every attention head its own query, key, and value projections. The heads can learn different relationships, then their outputs are concatenated and mixed. This is the attention mechanism introduced in Attention Is All You Need.

Grouped-query attention (GQA)

GQA keeps many query heads but groups them so that several query heads share one key and value head. It reduces key/value cache memory and bandwidth while retaining more query-head diversity than MQA. See GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.

Multi-query attention (MQA)

MQA gives every query head its own projection but shares a single key and value head across all queries. This greatly reduces key/value cache size and decoding bandwidth, at the cost of less independent key/value representations. It was introduced in Fast Transformer Decoding: One Write-Head is All You Need.

Multi-head latent attention (MLA)

MLA compresses each token's key and value information into a smaller latent vector, then reconstructs the head-specific representations when computing attention. The KV cache stores the latent vectors instead of every key and value head, reducing memory and decoding bandwidth while preserving more head diversity than MQA. It was introduced in DeepSeek-V2.

Sliding-window attention

Sliding-window attention (SWA) restricts each token to attending only to nearby tokens. In a causal decoder, query token $i$ attends to keys in the preceding window $[\max(0, i-w), i]$, where $w$ is the window size. The Longformer popularized a sliding-window attention design for long documents, combined with selected global attention links. Local attention predates Longformer, but Longformer made this sparse pattern a practical drop-in replacement for full attention.

In practice, SWA is often interleaved with full-attention layers. Using sliding windows everywhere would reduce the cost further, but a token could no longer directly retrieve information from distant parts of the prompt. Full-attention layers preserve these long-range retrieval paths while most layers keep the cheaper local pattern.

In forward, the attention kernel computes query-key scores and the weighted value sum only for the window edges. It does not materialize the full $n\times n$ attention matrix. In backward, the same sparse edges receive gradients: each query, key, and value is differentiated only through the windowed attention pairs that existed in forward. Stacking layers expands the receptive field, so information can travel roughly $Lw$ tokens across $L$ local-attention layers.

For a sequence of length $n$, full attention has $O(n^2)$ attention-score work per layer in both training forward and backward. SWA reduces this to $O(nw)$. During inference, prompt prefill falls from $O(n^2)$ to $O(nw)$. After the cache has reached the window size, each generated token has $O(w)$ attention work and key–value reads instead of $O(n)$, and the per-layer KV cache can remain $O(w)$ instead of growing as $O(n)$.

Linear attention

In practice, linear attention is often combined with full-attention layers. Its fixed-size state gives predictable memory and compute as the context grows, but compressing the past into that state can make it harder to retrieve an arbitrary distant token exactly. Full-attention layers restore direct long-range access while the linear layers keep most of the sequence processing cheaper.

Further reading:

• The Mamba-2 Model
• Linear attention
• DeltaNet
• Linear attention video
• Linear attention video

Attention sinks

Keeping only the most recent window can make a streaming language model unstable once the context exceeds the cache size. Attention sinks keep a small number $s$ of initial token KVs permanently, together with the latest $w$ token KVs. Initial tokens often receive high attention even when they are not semantically important, acting as a place for excess attention probability to go. Every new query can attend to the $s$ sink tokens and its local window, giving $O(w+s)$ decode work and bounded cache memory while preserving stable long-stream behavior.

PostNAS and Jet-Nemotron

Jet-Nemotron uses PostNAS to search for an efficient hybrid language-model architecture starting from a pretrained full-attention model. The search can choose which layers retain full attention, which use linear attention, and which new attention blocks to introduce, while also considering hardware-aware hyperparameters. This treats the placement of attention mechanisms as a search problem rather than a fixed architectural choice.

Encoder–decoder architecture

An encoder–decoder model runs an encoder over the input once, then generates output tokens with an autoregressive decoder that cross-attends to the encoder states. This split can improve inference speed for generation-heavy workloads. Suppose an encoder–decoder model has the same total parameter budget as a decoder-only model, but two-thirds of its parameters are in the encoder and one-third are in the decoder. The encoder work is paid once for the input, while the smaller decoder runs once for every generated token. The original T5 paper uses this encoder–decoder Transformer design for text-to-text tasks.

Generation is often memory-bandwidth-bound because each new token requires reading decoder weights and key–value cache data from HBM. A smaller decoder means less weight traffic on every decode step, so it can generate tokens faster. This is most useful when decoding dominates latency, such as long outputs with short or moderate inputs. The encoder prefill can be larger, but it is a one-time cost rather than work repeated for every generated token.

Weight tying

Language models usually have an input embedding matrix and a separate output projection, or language-model head. Both have one vector per vocabulary token. When the vocabulary is large relative to the rest of a small model, these two matrices can dominate parameter and memory use. Weight tying reuses the input embedding matrix $E$ for the output projection instead of learning a second vocabulary-sized matrix:

Here $h_t$ is the hidden state at position $t$, $b$ is an output bias, and $z_t$ contains one logit for every token. The logit for token $w$ is the dot product of $h_t$ with the embedding of $w$. Softmax therefore ranks tokens by their similarity to the hidden state, much like a similarity search in the token-embedding space. Tying also encourages tokens with similar embeddings to receive similar predictions.

In an autoregressive model, $h_t$ represents the token that should come next given the preceding context. Weight tying assumes that this prediction lives in the same semantic space as input token embeddings, so similarity-based prediction is meaningful. It halves the vocabulary-embedding parameters and was studied in Using the Output Embedding to Improve Language Models.

Mixture of experts (MoE)

An MoE layer contains several expert networks and a router. For each token, the router selects only a small number of experts, such as the top one or two, and the layer combines their outputs.

This lets the model increase its parameter count without running every parameter for every token. The main challenges are routing tokens evenly across experts and communicating tokens to the workers that own the selected experts.

Switch Transformer

Switch Transformer replaces only the FFN sublayer with an MoE. A dense FFN applies the same weights to every token:

The Switch layer instead has $E$ complete FFN experts:

A small router assigns each token a score for every expert, $r(x) = xW_r$, and sends the token to the highest-scoring expert. This is top-1 routing: only one expert runs for a token, so there is no averaging over several expert outputs. The model therefore has roughly $E$ times as many FFN parameters, while the FFN compute per token stays close to the dense baseline.

This can make MoE training faster than dense training for a target quality, though not every MoE step is faster. A sparse MoE activates only a few experts for each token, so its active compute can stay close to a dense feed-forward layer while its total parameter count and capacity are much larger. Different experts can specialize in different token patterns, giving more useful model capacity per unit of computation.

MoE also creates a smaller effective batch for each expert because every token is sent only to a subset of experts. Each expert therefore performs smaller matrix multiplications, which can be less efficient. A larger global batch is often needed to provide enough tokens to each expert and keep the same computations well saturated.

Some MoE layers also include a shared expert that processes every token alongside the selected routed experts. The shared expert can learn common features, while the routed experts specialize. For example, Qwen3.5-35B-A3B has 256 routed experts, selects 8 of them per token, and adds 1 shared expert. Thus 257 expert modules exist in total if the shared expert is counted, while 9 expert branches are active for each token.

Activation functions

Activation functions add nonlinearity to neural networks. Without them, stacking linear layers would still produce only one linear transformation.