[MaxGhenis]

Stacked on #1078 (l2-design-basis, which adds mass_parametrization="softmax"). Until #1078 merges, this PR's diff includes it; review b7b4c2935..HEAD. Once #1078 merges, main will be merged in here and the pins re-checked.

What changes

Under mass_parametrization="softmax", every Adam step ended on a loop: renormalize so exp(log_w) sums to the total, clamp to the ratio cap, repeat, for at most 32 rounds. This PR replaces that loop with the exact projection onto {w : sum(w) = total, w <= cap * d} in log space (_project_softmax_log_weights_ in solve.py). It is the Kullback-Leibler projection, log_w <- min(log_w + s, log cap), with one common shift s:

• How s is found. Active-set rounds, each one float64 pass over the records. With C the records already at their caps, the next shift solves sum(cap[C]) + exp(s) * sum(w[not C]) = total. That shift is a lower bound on the answer for every C, so C only grows, and a round that adds no record is exact.
• Worst case. An adversarial input can need one round per record. After 8 rounds, a sort over the remaining records finishes it, so the worst case is O(n log n).
• Receipt. options["iterate_selection_receipt"]["softmax_in_loop_max_cap_ratio"] records the largest ratio of a realized in-loop weight to its float64 cap, over every forward pass and the closing iterate, before the closing projection. It replaces softmax_cap_rounds_exhausted_epochs, which is always zero now. The ACS local-area tool carries the new key in its calibration summary and build manifest (last epoch batch, as before).
• Unchanged. The "projection" parametrization, the closing float64 projection and every default path. test_defaults_reproduce_the_pre_change_optimizer_bytes still passes against optimizer_bda72cb02.py, and so do the best-iterate oracle pins in test_solve.py.

Invariants (all property-tested)

test_softmax_cap_projection.py checks the projection for every input Hypothesis draws, 300 examples per property. The inputs include zero design weights (zero caps), a cap of exactly 1 (the caps sum to the total), and totals from 5% to 100% of the caps' sum.

1. Conservation and caps: sum(exp(P(x))) == total (rel 1e-12), and P(x) <= log cap exactly. A record with zero weight or a zero cap stays at zero.
2. Idempotence and shift invariance: P(P(x)) == P(x) and P(x + c) == P(x).
3. Feasible inputs: an input already within its caps after the shift moves by the shift alone.
4. KKT of the KL projection: every uncapped record moved by one common shift, and every capped record would pass its cap under it.
5. Differential against an independent reference. The result equals a float64 water-fill written from the definition. The active-set and sorted paths agree for 0, 1, 2 and 8 rounds. A constructed chain that needs 12 rounds exercises the sort fallback and matches the reference.
6. Float32, as the solver runs it: the clamp is exact in log space, and the realized total * softmax(log_w) sits within its float64 caps to 1e-5.

In test_l2_basis.py:

• Receipt bound: the receipt is <= 1 + 1e-5 for every softmax solve the property draws. That includes caps of 1, 1.001 and 1.05, where the cap binds on nearly every record; at a cap of 1 the solve returns the design weights.
• No new receipt under projection: projection solves record no softmax receipt.
• CLARABEL path tests unchanged: they pass with the new step.

Weights are compared in weight units (atol 1e-12 × total), not log units. When the caps nearly meet the total, the uncapped records hold only a small remainder, and their float64 rounding is a few ulps of the total. That error is large in log space and small in weight space.

Full-scale measurement (Modal, off the shared Mac)

experiments/us-acs-local-l2-basis-20260928/exact_cap_modal.py reran the ACS local release surface on Modal: 1,588,854 households, 4,459 targets, a 5x cap and 800 epochs. It used two configurations whose earlier runs ran out of rounds. Each ran under the new step and under #1078's loop (copied into the script), with the largest in-loop ratio recorded at every forward pass (results/exact_cap.md):

Configuration	Cap step	Largest in-loop weight / cap	Steps out of rounds	Cap step	Final loss	National ESS
share 0.5, chi-square λ 0.03	exact projection	1 + 9.0e-6	0 of 800	12.7 ms	0.018991	21,834
share 0.5, chi-square λ 0.03	32-round loop	1 + 1.2e-5	713 of 800	53.8 ms	0.019169	21,835
share 0.9, no penalty	exact projection	1 + 9.0e-6	0 of 800	22.5 ms	0.031293	93,609
share 0.9, no penalty	32-round loop	1 + 8.7e-6	705 of 800	48.7 ms	0.031365	93,598

What this shows:

• The documented limitation was smaller than feared. Chi-square design-weight penalty and softmax mass parametrization for calibration; ACS local ESS frontier #1078's docs could not say how far past the cap the loop left the weights. The answer is a float32-sized amount, about the same as the exact step's.
• Both residuals come from the solve's float32 softmax. Its normalizer missed the total by up to 8.2e-6. With a float64 normalizer, the projected vectors sit within 1 + 4.9e-7 of their caps (exact_cap_float32.py, results/exact_cap_float32.json).
• The loop ran out because of float32 too. Its shift was log(total) in float64 minus a float32 logsumexp. That bottoms out at log(total)'s float32 rounding residual (7.8e-8, measured exactly). Each round, that shift lifts one to three capped records whose log caps are below 2 in magnitude one float32 step over their caps. On vectors the projection had already made exact, the loop still ran out in 10 of 12 trials.
• What the change buys:
  • the step is exact and never runs out (one to three rounds on every step at full scale);
  • it is 2.2-4.2x cheaper;
  • the in-loop overshoot is now recorded rather than unknown.
• Final metrics. ESS agreed within 0.1%. At λ = 0.03, the loop run reproduced the grid's run exactly (0.019169), and the earlier rerun gave 0.019104. The projection's 0.018991 is 0.6-0.9% below those three, from one run, so not established.

docs/calibration-l2-basis.md replaces the "known limitation" section with this account. The experiment README adds it. Projection remains the recommended parametrization at that scale: softmax's training loss at λ = 0.03 is still 5% above projection's.

Re-pins (solve.py's bytes are attested)

• EXPECTED_HASHES: seed_protocol be39eb6b -> 5795dc07, seed_map 085d8d39 -> 292ad19c. tools/spec_engine_coverage.py --check reports 41/41.
• US spec_sha256 e0b757ce -> edcd6475, in us-f0-coverage.json and test_us_multispine_pool_tool.py.
• Loader golden fee5893a -> 5135dd79.
• Calibrate parity pins, via tools/graph_parity_repin.py calibrate; the direct-call bytes are unchanged.

Tested locally (targeted, on the shared host)

• Calibrate tests: test_softmax_cap_projection.py (11 passed) and test_l2_basis.py (all 51 passed, including the CLARABEL path tests and the byte pin).
• Tool and oracle tests: the ACS tool's penalty and settings tests, and the best-iterate oracle tests in test_solve.py.
• Pins: the loader golden and the graph parity pins (12 passed).
• Not run locally: the engine jobs (engine-us), which CI runs.

🤖 Generated with Claude Code