Activations, Optimizers & Systems

GELU and SwiGLU, LayerNorm and RMSNorm, RoPE, Adam. The unglamorous machinery that makes a deep stack trainable at all.

Published: 2 August 2026
Reading time: 8 min read
Figures: 4 figures
Equations: 9 equations

Activation Functions: ReLU → GELU → SwiGLU

Activation functions introduce nonlinearity — without them, any stack of linear layers collapses to a single linear transformation.

ReLU f(x) = max(0,x): simple, fast, but the 'dying ReLU' problem — neurons whose inputs are always negative output 0 with exactly zero gradient, permanently stopping learning.

GELU (Hendrycks & Gimpel, 2016): f(x) = x·Φ(x) where Φ is the standard normal CDF. This smoothly gates each input by how likely it is to be positive. Unlike ReLU's hard cutoff, GELU gives small non-zero gradients for negative inputs. Approximated by a tanh formula for efficient computation. Used by BERT, GPT-2/3.

SwiGLU (Shazeer, 2020): a gated architecture. FFN(x) = Swish(xW₁)⊙(xW₂), where Swish(x) = x·σ(x). The second branch acts as an input-dependent gate. Because this introduces a third projection matrix, the inner dimension is reduced to (2/3)×4d_model to keep total FLOPs constant. Empirically the best activation for Transformers. Now default in LLaMA, PaLM, Mistral, Qwen, DeepSeek.

ReLU

ReLU (x) = max (0, x)

GELU (approximation)

GELU (x) \approx 0.5 x (1 + tanh (\frac{2}{π} (x + 0.044715 x^{3})))

SwiGLU

SwiGLU (x, W_{1}, W_{2}) = Swish x W_{1} \cdot σ (x W_{1}) ⊙ (x W_{2})

Activation Functions: ReLU → GELU → SwiGLU

python

import torch
import torch.nn.functional as F

def swiglu(x, W, V):
    # x: (batch, dim)
    # W, V: linear projections
    gate = F.silu(x @ W)  # SiLU is Swish(x)
    val = x @ V
    return gate * val

Normalization: LayerNorm vs. RMSNorm

Without normalization, the distribution of each layer's inputs shifts as earlier layer weights update — 'internal covariate shift' — causing training instability and requiring very low learning rates.

LayerNorm (Ba et al., 2016): for each example independently, normalize across the feature dimension using mean and variance, then apply learnable scale γ and shift β. Different from BatchNorm (normalizes across the batch), which is inappropriate for variable-length sequences.

Pre-Norm vs Post-Norm: the original Transformer applied LN after the residual: x = LN(x + sublayer(x)). Modern models (GPT-2 onward) apply LN before: x = x + sublayer(LN(x)). Pre-Norm preserves a clean residual stream, enabling much deeper and more stable networks.

RMSNorm (Zhang & Sennrich, 2019): drops the mean-centering step — normalizes only by the root mean square of activations. ~15% cheaper to compute, empirically matches LayerNorm quality. Now standard in LLaMA, Mistral, Gemma, Qwen, DeepSeek.

LayerNorm

LN (x) = γ \cdot \frac{x - μ}{σ ^{2} + ε} + β

RMSNorm

RMSNorm (x) = \frac{x}{RMS ( x )} \cdot γ, RMS (x) = \frac{1}{d} i \sum x_{i}^{2}

Pre-Norm (modern)

x \leftarrow x + sublayer (RMSNorm (x))

Normalization: LayerNorm vs. RMSNorm

python

import torch

class RMSNorm(torch.nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = torch.nn.Parameter(torch.ones(dim))

    def forward(self, x):
        # Root Mean Square
        rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps)
        return (x / rms) * self.weight

RoPE — Rotary Position Embedding

RoPE (Su et al., 2021) is the dominant positional encoding in modern LLMs. Core idea: encode position by rotating Q and K vectors in 2D subspaces.

For each dimension pair (2i, 2i+1), apply a 2D rotation matrix by angle m·θᵢ, where m is the token position and θᵢ = 10000^{−2i/d} is a frequency (geometric progression, like sinusoidal PE).

The crucial mathematical property: the dot product q_m·k_n after applying their rotations equals f(m−n) — it depends only on the relative offset, not the absolute positions. Relative position emerges naturally from absolute rotations.

Advantages: parameter-free (no learned table); works with linear attention; extrapolates more gracefully than learned absolute embeddings; compatible with standard self-attention kernels.

Used by: LLaMA (all versions), Mistral, Qwen, Falcon, Gemma, DeepSeek.

YaRN (Peng et al., 2023): extends RoPE beyond training context by interpolating rotary frequencies — ~2.5× fewer fine-tuning steps than position interpolation (PI).

2D rotation at position m

[q_{2 i}^{'} q_{2 i + 1}^{'}] = [cos m θ_{i} sin m θ_{i} - sin m θ_{i} cos m θ_{i}] [q_{2 i} q_{2 i + 1}]

Frequency schedule

θ_{i} = 1000 0^{- 2 i / d}

Relative distance property

q_{m}^{⊤} k_{n} = f (m - n) (only depends on relative offset)

RoPE — Rotary Position Embedding

python

def apply_rope(q, k, positions):
    # Rotate adjacent pairs of dimensions
    q_rotated = q * cos(positions) + rotate_half(q) * sin(positions)
    k_rotated = k * cos(positions) + rotate_half(k) * sin(positions)
    return q_rotated, k_rotated

Distributed Training Dimensions

Training frontier models requires splitting the workload across thousands of GPUs.

• Data Parallelism (DP/FSDP): Copy the model weights to every GPU, split the data. FSDP shards the optimizer states and weights to save memory. • Tensor Parallelism (TP): Split individual matrix multiplications across GPUs (e.g., Q, K, V projections). Requires extremely high-bandwidth interconnects (NVLink) because of all-reduce operations. • Pipeline Parallelism (PP): Split the model by layers across GPUs (GPU 1 computes layers 1-4, then sends activations to GPU 2 for layers 5-8). • Expert Parallelism (EP): For MoE models, different GPUs host different experts. Tokens are routed across the network to the correct GPU.

Modern Optimizers (AdamW to Muon)

AdamW has been the default optimizer since 2017, combining adaptive learning rates (using momentum and variance) with decoupled weight decay.

However, AdamW scales memory poorly because it keeps two state variables (moment 1 and 2) per parameter.

Newer optimizers are challenging this: • Sophia (2023): A second-order optimizer using a light-weight diagonal Hessian estimate. Updates clip gradients dynamically. Claims 2x faster convergence than Adam. • Muon (2024): Used in DeepSeek training. It applies orthogonalization (Newton-Schulz iteration) to the gradients of weight matrices, effectively preconditiong them without storing large momentum tensors. Extremely memory efficient.

FlashAttention Internals

While conceptually solving the memory bottleneck, FlashAttention is really a masterclass in hardware-aware programming. GPUs have a massive but slow memory pool (HBM) and a tiny but incredibly fast memory pool (SRAM).

Standard attention writes intermediate matrices (like QK^T) to HBM, then reads them back to apply Softmax. This IO bottleneck dominates the runtime.

FlashAttention aggressively tiles the Q, K, and V blocks to fit exactly into SRAM. It computes the attention incrementally (online softmax) and only writes the final output back to HBM. FlashAttention-3 goes further, exploiting the Hopper architecture's asynchronous DMA and Tensor Cores to overlap data movement with math.

FlashAttention Internals