AI Engineer Exhaustive
MediumPlusgpt-oss

Attention Sinks

gpt-oss

Medium

Learned per-head scalar added inside the softmax denominator, giving each attention head an escape valve so weights no longer sum to 1.

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Section 2.2 of the gpt-oss model card states that each attention head has a learned bias in the softmax denominator, “similar to off-by-one attention and attention sinks,” so a head can pay no attention to any token. That is the whole architectural claim. The paper does not write an equation. The reference implementation in gpt_oss/torch/model.py does: it stores one scalar per query head and concatenates it as an extra logit before softmax.

Ordinary softmax on a length-(T) score row is a distribution over the (T) keys. The weights always sum to one, so every query is forced to mix some value vector even when none of the visible keys are useful (for example a sliding-window layer whose band is empty of relevant context, or a head that has already written its information into the residual). A learned sink logit (\sigma_h) adds a dummy class that absorbs residual mass. After softmax, that dummy weight is dropped. The remaining weights over real keys then sum to (1 - p_{\text{sink}} < 1).

The mechanism is per head, not per token or per key. All query positions of head (h) share the same (\sigma_h). Grouped-query attention does not share sinks across a KV group: there are 64 sink scalars, one per query head.

How it works

Let (s_{h,t,j}) be the already-scaled, already-masked score of query (t) against key (j) on head (h). The official sdpa routine builds

\tilde{s}_{h,t} = \bigl[s_{h,t,0},\;\ldots,\;s_{h,t,T-1},\;\sigma_h\bigr] \in \mathbb{R}^{T+1}

and takes (\mathrm{softmax}) over the last axis. The first (T) entries become the attention weights used against (V); the last entry is discarded:

\alpha_{h,t,j} = \frac{e^{s_{h,t,j}}}{e^{\sigma_h} + \sum_{k=0}^{T-1} e^{s_{h,t,k}}},\qquad j=0,\ldots,T-1.

That is exactly a learned additive term in the denominator. If (\sigma_h \to +\infty), every (\alpha_{h,t,j}\to 0) and the head writes a zero residual (before the output projection). If (\sigma_h \to -\infty), the extra term vanishes and the head recovers ordinary softmax.

In the reference code the sink tensor S has shape [num_attention_heads]. After Q is viewed as [T, n_kv, q_mult, d], S is reshaped to [n_kv, q_mult, 1, 1] and expanded to [n_kv, q_mult, T, 1], then concatenated on the key axis of QK. Masked (future or out-of-window) scores stay at (-\infty); the sink is never masked, so it remains a valid escape even when the window is tiny.

The paper cites Miller’s “attention is off by one” (a fixed (+1) in the denominator) and Xiao et al.’s streaming sinks (a dedicated sink token). gpt-oss is closer to the first: there is no extra key/value vector, only a logit. The official graph does not add (\sigma_h) to every score; it appends it. Adding it to every entry would cancel in softmax and do nothing.

Shapes for the 120b/20b configs: 64 query heads, so 64 sink parameters in bfloat16. They are independent of sequence length and of the 8 KV heads.

Official code

gpt_oss/torch/model.py: AttentionBlock constructs self.sinks as an nn.Parameter of length config.num_attention_heads. The function sdpa(Q, K, V, S, sm_scale, sliding_window) concatenates S onto the score tensor, softmaxes, then slices W[..., :-1] before the value mix. The educational PyTorch path is the source of truth; the Triton and Metal backends implement the same extra logit.

Watch-outs

Sources