Gated Attention

Context: Gating Mechanisms

Gating mechanisms have a long history in neural network architecture as a way to regulate information flow and improve the propogation of gradients $(\nabla )$.

Early architectures like LSTMs, Highway Networks, and GRUs pioneered gating to control flow across layers / time steps.

More recent works like State-Space Models (SSMs) and linear attention variants often also apply gating to modulate token-mixer outputs.

Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free

The paper systematically investigates augmenting standard Softmax attention with gating at various positions $(G_1,G_2,G_3,G_4,G_5)$.

Optimal Position $\left(G_1\right)$: Applying a head-specific sigmoid gate after the Scaled Dot-Product $\left(\mathrm{SDPA}\right)$ yields the most significant improvements.

Mathematics: Formalizing Gated Attention

The gating mechanism is formalized as a dynamic filter applied to a target input $Y$ using a gating input $X$:

$Y^{\prime }=g(Y,X,W_{\theta },\sigma )=Y\odot \sigma (X{W}_{\theta })$

where $Y$ is the input to be modulated (e.g., $\mathrm{SDPA})$, $X$ is the hidden states after pre-normalization used to compute gating scores, and $\sigma$ is the activation function, typically ${\mathrm{sigmoid}}_{}$. The effectiveness is attributed to two factors:

1. Non-linearity: This increases the expressiveness of the low-rank linear transformation formed by the value $\left(W_V\right)$ and the dense $\left(W_O\right)$ projections.

2. Input-Dependent Sparsity: Gating scores concentrated near zero introduce sparsity, filtering irrelevant context.

Impact: Performance and Scaling

They see a reduction of over $0.2$ PPL in $15$ billion parameter MoE models using $\mathrm{SDPA}$ Elementwise $G_1$.

Training Stability: The gated architecture nearly eliminates loss spikes, allowing for larger learning rates and better scaling. Gating allows dense models to converge at higher learning rates (e.g., doubled to $8.0\times {10}^{-3}$) where standard baselines would otherwise diverge.

Attention-Sink-Free: Sparse gating scores effectively eliminate the "attention sink" phenomenon, where models over-allocate attention to the first token. In baseline models, the first token receives $~46.7\text{\%}$ of attention; gating reduces this to $~4.8\text{\%}$ of attention.

Context Extension: SDPA output gating facilitates adaptation to sequence length increases by replacing fixed attention-sink patterns with input-dependent filtering. Gated models demonstrate superior length generalization, outperforming baselines significantly at $64$k and $128$k context lengths on the RULER benchmark.

Questions

• N/A