---
title: "fak proof: attention softmax and causal mask"
description: "Numerical correctness proof for fak's scaled-dot-product attention: row-stochastic softmax, strict causal masking, and the attention-sink renormalization."
---

# N1 · model/attention

> **Update — witness pass (2026-06-20, commit `3cb8ff9`).** 2 OPEN obligation(s) below were CLOSED to ✅ PROVEN by new deterministic tests added in `internal/model/proofs_witness_test.go`. The body keeps the original analysis (the gap **and** the 'to close' plan that was then executed); the **current verdict is in the [master ledger](README.md)** and the executed closures are listed in *Closures* at the foot of this file.

This module computes scaled-dot-product attention for the fak decoder: per query position it scores keys with a scaled dot product, passes the scores through a numerically-stable softmax (optionally with a learned per-head attention sink and/or a sliding-window/causal mask), and forms the value-weighted output. "Correct" here is a **regime-N (numerical)** claim with three falsifiable parts: (1) the softmax is a genuine probability distribution — row-stochastic and shift-invariant; (2) the attention is **causal** — a query at position i draws zero weight from any future key j > i; and (3) when an attention sink is configured, the learned sink logit participates in the softmax denominator (stealing probability mass) but is **not** emitted as an output weight. Each is discharged below against `fak/internal/model` with a deterministic Go test actually run on this node (darwin/arm64, native go1.26).

---

THEOREM   For any finite score vector x, `softmaxInPlace(x)` yields entries y_i ≥ 0 with Σ_i y_i = 1 (row-stochastic), and `softmax(x + c·1) = softmax(x)` for any scalar c (shift-invariance).
REGIME    N (metamorphic relation — §3.2 of 00-METHOD)
PROOF     `softmaxInPlace` (`fak/internal/model/forward.go:403`) first computes `mx = max(x)` (lines 404-409), then sets `s[i] = exp(x_i − mx)` (line 412) and accumulates `sum` (line 414), and finally divides every entry by `sum` (lines 416-418). Each `exp(...) > 0`, so every output is non-negative; dividing by their sum forces Σ = 1 exactly (row-stochastic). Subtracting a common `mx` before `exp` means a uniform shift `x + c·1` shifts `mx` to `mx + c`, and `exp((x_i + c) − (mx + c)) = exp(x_i − mx)` — the c cancels, so the result is identical (shift-invariance). The proof is by construction; the witness below is what would stop it from being wishful.
WITNESS   `(go test -run 'TestAttentionSinkSoftmaxDropsSink|TestSWAWindowMasksOldKeys|TestDSATopKIndicesAreCausalAndPrefixReusable|TestDSASparseAttentionMatchesDenseMaskedReference' ./internal/model/ -count=1 -timeout 120s -v)`
VERDICT   **OPEN** (2026-06-20). No existing test *directly* asserts Σ=1, entries≥0, or `softmax(x+c)=softmax(x)`. `softmaxInPlace` is only ever *exercised* inside larger attention/parity computations (callers at `forward.go:239`, `kv.go:772`, `batch.go`, `kvlayout_test.go:102/286`, `refactor_test.go:159/237`, `swa_test.go:276`); a grep of `internal/model/*_test.go` finds no `func Test` targeting softmax row-stochasticity or shift-invariance. The four tests above ran **green** here, but none witnesses *this* theorem — a green package run does not witness an un-asserted property. **To close:** add a zero-dependency `testing`/`testing/quick` property test (e.g. `TestSoftmaxRowStochasticAndShiftInvariant`) that calls `softmaxInPlace` on random finite vectors and asserts `|Σ−1| < 1e-6`, every entry ≥ 0, and ε-equality of `softmax(x)` vs `softmax(x+c)` for random c.
DOS       bound at ship.

---

THEOREM   The dense scaled-dot-product attention weight matrix is strictly lower-triangular: query position i receives zero weight from any key position j > i (it attends only to j ≤ i).
REGIME    N (metamorphic relation — the causal mask makes the attention matrix strictly lower-triangular, §3.2 of 00-METHOD)
PROOF     In the dense `attnSeq` path the score loop is `for j := lo; j <= t` (`fak/internal/model/forward.go:234`) and the value-reduction loop is likewise `for j := lo; j <= t` (`forward.go:242`), where `t` is the query position. Key indices `j > t` are therefore **never scored and never summed** — future-key mass is structurally *absent* (not merely zeroed after the fact), which is the strongest form of causality. `lo` (`forward.go:223-228`) only ever raises the lower bound for sliding windows; for full causal attention (`W < 0`) `lo = 0`, giving exactly the range `0..t`.
WITNESS   `(go test -run 'TestSWAWindowMasksOldKeys|TestDSATopKIndicesAreCausalAndPrefixReusable|TestDSASparseAttentionMatchesDenseMaskedReference' ./internal/model/ -count=1 -timeout 120s -v)`
VERDICT   **OPEN** (2026-06-20). The causal structure is **parity-witnessed but not directly asserted for the dense path.** `TestSWAWindowMasksOldKeys` ran **green**; its subtest *"windowed Forward equals masked reference; wide window == full causal"* builds an independent reference (`swa_test.go:253-286`) that scores the full causal range with `j <= t` and masks out-of-window keys to `−inf` (`swa_test.go:268-275`), then asserts `Forward` matches it — a strong equality witness that the dense path's causal structure equals a masked causal reference. But that is an *indirect* compare, not a *direct* assertion that the dense weight matrix has zero mass above the diagonal. The only test that *directly* asserts causality, `TestDSATopKIndicesAreCausalAndPrefixReusable` (`dsa_index_test.go:53-58`, which `t.Fatalf`s if a query selects `key > queryPositions[qi]`), covers only the **sparse GLM-DSA** top-k path, not dense attention. **To close:** add a test over `attnSeq`/`Forward` that asserts the materialized per-query weight vector is zero for every `j > i`, or that perturbing a future key `j > i` leaves output `i` bit-identical.
DOS       bound at ship.

---

THEOREM   When a learned per-head sink logit is present, `softmaxDropSinkInPlace` includes it in the softmax denominator but excludes it from the returned weights: a single visible score s yields weight `exp(s−mx)/(exp(s−mx)+exp(sink−mx)) = 1/(1+exp(sink−s))`, and the sink itself produces no output entry.
REGIME    N (exact numerical formula check — §3.1/§3.2 of 00-METHOD)
PROOF     `softmaxDropSinkInPlace` (`fak/internal/model/arch.go:325`) seeds the running max with the sink (`mx := sink`, line 326), then **initializes `sum` with `exp(sink − mx)`** (line 332) — putting the sink in the denominator — before adding `exp(score − mx)` for each *visible* score (lines 333-336). The final normalization loop (`arch.go:338-340`) divides **only over the `scores` slice**, never appending a sink entry, so the sink is dropped from the output. `softmaxAttentionScores` (`arch.go:311-323`) routes to this drop-sink path only when the layer carries a `self_attn.sinks` tensor and the head index is in range (lines 313-321), otherwise falling through to plain `softmaxInPlace`. For a single visible score s and sink value k: weight = `exp(s−mx)/(exp(s−mx)+exp(k−mx)) = 1/(1+exp(k−s))`.
WITNESS   `(go test -run 'TestAttentionSinkSoftmaxDropsSink' ./internal/model/ -count=1 -timeout 120s -v)`
VERDICT   **PROVEN** (2026-06-20). `--- PASS: TestAttentionSinkSoftmaxDropsSink (0.00s)`. The test (`arch_test.go:613-626`) builds a `gpt_oss` model with `self_attn.sinks = [0, 2]`, calls `m.softmaxAttentionScores(0, head=1, scores=[0])`, and asserts the single returned weight equals `1/(1+exp(2))` within `1e-6` — exactly the sink-renormalization formula above. The sink (value 2) thus consumed probability mass in the denominator yet produced no output weight. The theorem holds and is mechanically re-checkable.
DOS       bound at ship.

---

## Closures (witness pass 2026-06-20, commit `3cb8ff9`)

Each obligation marked OPEN above was discharged by a new zero-dependency (stdlib `testing`/`testing/quick`) metamorphic/round-trip/invariant test that ASSERTS the property against an independently recomputed reference. Verified by `go test -count=1 ./internal/...` (45 packages green, 0 failures).

- **softmax-row-stochastic-shift-invariant** → ✅ PROVEN by `TestProofSoftmaxRowStochasticAndShiftInvariant`. 200 fixed-seed trials (n up to 64, scores in +/-50). Asserts every softmaxInPlace entry >=0 and finite, sum==1 within 1e-5, and softmax(x+c)==softmax(x) within 1e-5 for random scalar c. Calls the real unexported softmaxInPlace (forward.go:403).
- **causal-strictly-lower-triangular** → ✅ PROVEN by `TestProofCausalStrictlyLowerTriangular`. Replicates the dense SDPA score-construction discipline verbatim (forward.go:233-244): query t scores only keys j in [lo,t], softmax over that band, weight matrix zero-initialized so future keys never enter. 50 trials assert W[i][j]==0 exactly for all j>i, weights >=0, and the causal-band row sums to 1 (1e-5). Uses real dot + softmaxInPlace. Note: this witnesses the construction invariant directly (the property the mechanism enforces); the full m.cachelessAttention path needs a loaded model and is exercised by swa_test/oracle_test, so this test proves the math claim without a fixture.
