Softmax-Free Attention in JAX/Flax NNX

The explainer swaps attention’s exp(q·k) for the Yat kernel and reports what a parameter-matched contest measures. This is the implementation: one attention module with two score branches, the training harness that holds everything else fixed, and the telemetry. Everything below is scripts/yat_attention.py, the script both Kaggle runs executed (quality bundle kgl_blog-yatattn-v1, telemetry bundle kgl_blog-yatattn-telem); code blocks are quoted from it.

One module, two compatibility functions#

The entire experimental contrast lives inside Attention.__call__. Same projections, same heads, same causal mask, same value aggregation; the branch chooses only how a query and key become a score, and how scores become weights:

if s.variant == "softmax":
    logits = (q @ k.transpose(0, 1, 3, 2)) / math.sqrt(s.dh)
    logits = jnp.where(mask, logits, -jnp.inf)
    w = jax.nn.softmax(logits, axis=-1)
else:
    dots = q @ k.transpose(0, 1, 3, 2)                       # (B,h,T,T)
    d2 = (jnp.sum(q * q, -1, keepdims=True)
          + jnp.sum(k * k, -1, keepdims=True).swapaxes(-1, -2)
          - 2 * dots)
    b = jax.nn.softplus(s.b_raw.value)[None, :, None, None]    # per head
    eps = jax.nn.softplus(s.eps_raw.value)[None, :, None, None]
    kappa = ((dots + b) ** 2) / (jnp.maximum(d2, 0.0) + eps)   # >= 0, Mercer
    kappa = jnp.where(mask, kappa, 0.0)
    mass = kappa.sum(-1, keepdims=True)                        # (B,h,T,1)
    w = kappa / (mass + 1e-9)

The kernel scalars are stored raw and passed through a softplus, so admissibility ($b, \varepsilon > 0$) is structural rather than clipped, one pair per head, initialized at zero so the effective values start at $\log 2$:

if variant == "yat_b":
    s.b_raw = nnx.Param(jnp.zeros((heads,)))
    s.eps_raw = nnx.Param(jnp.zeros((heads,)))

Read the kernel branch against the softmax branch and notice what is missing. There is no exponential, so there is no max-subtraction trick and no shift gauge. The causal mask is 0.0, not -inf, because zero kernel mass is exclusion. And mass, the row sum that softmax divides away inside its own definition, is a first-class tensor here, kept for telemetry. The pairwise distance is assembled from the same dots the score already needed, so the kernel branch costs one extra elementwise pass over the attention matrix, not an extra matmul.

Everything around this branch is deliberately boring: the same projections, the same pre-norm GPT skeleton, the same AdamW, the same batches in the same order. Softmax counts 2,719,169 parameters and the kernel 2,719,217; the difference is exactly the forty-eight learned scalars.

The contest#

Three seeds per variant, 12,000 steps, each at its swept learning rate (3e-4 softmax, 3e-3 kernel), best validation loss on held-out Shakespeare:

The telemetry that killed a claim#

The run logs the maximum attention score each model produces on held-out batches. We expected to caption this “bounded kernel, roaming logits”. The data said otherwise, and the companion is where the numbers live:

The kernel’s scores are five orders of magnitude larger, because $\kappa(q,q) = \lVert q \rVert^4 / \varepsilon$ and training grows the projection norms. The lesson the explainer draws is the honest one: boundedness-for-bounded-inputs is not boundedness-under-training, and what actually matters is that a ratio of polynomials digests $4 \times 10^5$ without a max trick, while $e^{34.5}$ does not exist in float32 at all.

The mass channel’s naive test is measured with a rank statistic, mass against next-token correctness:

order = np.argsort(m)
ranks = np.empty_like(order, dtype=float); ranks[order] = np.arange(len(m))
pos = ranks[c == 1]
auroc = (pos.sum() - len(pos) * (len(pos) - 1) / 2) / (len(pos) * (c == 0).sum() + 1e-9)

AUROC 0.509, 0.501, 0.508 across seeds. Null, reported as such.

The learning-rate audit#

The Yat training recipe says kernel attention wants roughly three times the softmax learning rate, and the original contest ran both variants at softmax’s 3e-4. The audit sweep (bundle kgl_blog-yatattn-lrsweep, one seed per cell, full 12,000 steps each):

Three seeds re-run at the kernel’s best rate (bundle kgl_blog-yatattn-fair) give 1.5163 ± 0.0023, within run-to-run noise of the published 1.5131 ± 0.0071, so the quality table stands.

The b = 0 ablation row is not just a smaller kernel, it is a different mathematical object: the explainer walks through why the universality theorem needs the bias. Empirically the ablation costs a run-noise-sized amount of quality and, more decisively, un-tames the score scale: 443,508 maximum trained score against the full kernel’s 61,221.

The learned scalars themselves are exported by the telemetry run (bundle kgl_blog-yatattn-b-telem), and they tell a per-head specialization story: b wanders from its log 2 init across [0.22, 3.40], and eps splits, most heads driving it toward zero (sharp, local kernels) while two late-layer heads grow it to 6.4 and 12.7 (smooth, near-global ones). The entropy statistics on the trained maps barely move against the ablation (mean 0.73 vs softmax’s 0.48), so the diffuseness mechanism the explainer describes is a property of the polynomial ratio, not of the frozen scalars. The column shapes are the finding: softmax needs its rate found (and dies at 1e-2), while the kernel’s entire column sits in a 1 percent band. Per-variant LR overrides and the sweep mode are in scripts/yat_attention.py (SWEEP_LRS, LR_YAT, LR_SOFTMAX env knobs).

The routing, watched#

The telemetry run snapshots both models’ full attention tensors on one fixed validation window at four checkpoints. Two honest ways to look at them: through depth, and through training.

Watching a single head read is where the diffuseness difference stops being a statistic. One query position advancing through the window, its full attention row at every position:

Every number above is from scripts/yat_attention.py on Kaggle (bundles kgl_blog-yatattn-v1 and kgl_blog-yatattn-telem); figures are rendered from the bundles by scripts/render_yatattn_gifs.py. The linear-time road, positive random features for exactly this kernel, is SLAY (arXiv 2602.04915).

@misc{bouhsine2026attentionisacompatibilitykerneljaxflaxnnx,
  author       = {Bouhsine, Taha},
  title        = {Softmax-Free Attention in JAX/Flax NNX},
  year         = {2026},
  month        = {jul},
  howpublished = {\url{https://tahabouhsine.com/blog/attention-is-a-compatibility-kernel-jax-flax-nnx/}},
  note         = {Blog post, Records of the !mmortal Data Scientist}
}