Project the softmax parametrization exactly onto the capped simplex each step

changelog.d/calibrate-chi-square-l2-basis.added.md

@@ -1 +1 @@
-Add two opt-in options to Adam calibration and the ACS local-area tool. `l2_basis="chi_square"` (`--l2-basis chi_square`) makes the `l2_lambda` penalty GREG's design-weighted chi-square distance `sum(d * (w / d - 1) ** 2) / sum(d)`, which pulls toward the design weights themselves; the historical record-weighted `mean((w / d) ** 2)`, whose mass-conserved optimum is `w ∝ d ** 2`, stays the default. `mass_parametrization="softmax"` (`--mass-parametrization softmax`) conserves mass through `w = total * softmax(log_w)`; the historical per-step uniform shift, still the default, can stall away from the constrained optimum when every record's gradient shares a sign (on the ACS local release it does not, and softmax's per-step cap rounds run out on most epochs there, which the solve counts). Defaults reproduce the previous solver's bytes. Results record both settings and `CalibrationResult.chi_square_distance`; the tool records them, with the softmax cap-exhaustion count, in its calibration summary and build manifest, names all three penalty flags in the manifest's refresh recipe, and refuses to resume a checkpoint solved under other settings. See docs/calibration-l2-basis.md.
+Add two opt-in options to Adam calibration and the ACS local-area tool. `l2_basis="chi_square"` (`--l2-basis chi_square`) makes the `l2_lambda` penalty GREG's design-weighted chi-square distance `sum(d * (w / d - 1) ** 2) / sum(d)`, which pulls toward the design weights themselves; the historical record-weighted `mean((w / d) ** 2)`, whose mass-conserved optimum is `w ∝ d ** 2`, stays the default. `mass_parametrization="softmax"` (`--mass-parametrization softmax`) conserves mass through `w = total * softmax(log_w)`; the historical per-step uniform shift, still the default, can stall away from the constrained optimum when every record's gradient shares a sign (on the ACS local release it does not). Defaults reproduce the previous solver's bytes. Results record both settings and `CalibrationResult.chi_square_distance`; the tool records them in its calibration summary and build manifest, names all three penalty flags in the manifest's refresh recipe, and refuses to resume a checkpoint solved under other settings. See docs/calibration-l2-basis.md.

changelog.d/calibrate-softmax-exact-cap.changed.md

@@ -0,0 +1 @@
+`mass_parametrization="softmax"` now projects the log-weights exactly onto the capped simplex after every Adam step: one common shift, then a clamp to each cap. Active-set rounds find the shift in float64, and a sort finishes it in log space if they run long or the uncapped weights underflow float64. It replaces a renormalize/clamp loop of at most 32 rounds, which ran out on most epochs of the ACS local release because of float32 rounding (the overshoot it left was float32-sized, measured in `experiments/us-acs-local-l2-basis-20260928/results/exact_cap.md`). Softmax solves record `options["iterate_selection_receipt"]["softmax_in_loop_max_cap_ratio"]`, the largest in-loop weight over its cap, in place of the exhaustion count, and the ACS local-area tool carries it in its calibration summary and build manifest. The default `"projection"` parametrization is unchanged.

docs/calibration-l2-basis.md

@@ -95,29 +95,84 @@ The total holds by construction, and autograd hands Adam the gradient with its
 component along the constraint already removed. For a smooth objective that
 reduced gradient vanishes at the constrained optimum; the capped-MAPE loss has
 kinks, so Adam still oscillates there on the scale of `lr`, as it does under
-`mass="free"`. The ratio cap remains a per-step clamp, alternated with the
-softmax-invariant renormalization for up to 32 rounds, plus the closing
-float64 projection that both parametrizations share. It requires
-`mass="conserve"` and no L0 gates.
-
-**Known limitation: the cap rounds run out at production scale.** Each
-renormalization lifts the clamped records back over the cap by a shrinking
-amount, so the rounds converge geometrically but need not reach float32
-exactness in 32. A step that runs out ends on a clamp, and its next forward
-pass optimizes `total·softmax(log_w)` past the cap, by an amount that is not
-recorded.
-`options["iterate_selection_receipt"]["softmax_cap_rounds_exhausted_epochs"]`
-counts those steps. On small problems it is rare. On the ACS local release it
-is the normal state: 165-400 of the last 400 epochs in every run that records
-it, including 400 at the share-0.5, `l2_lambda = 0.03` configuration on the
-full surface and on one fold, and 400 in every share-0.9 run. Only the
-closing projection makes the returned vector exact. The returned loss stays
-within 0.3% of the last trajectory loss in every run that records the count,
-but projection runs, which have no cap loop, show the same gap (0.01-0.33%), so
-that says nothing about the overshoot; its size is not measured. An
-exact capped-softmax step (water-filling the excess onto the uncapped records)
-would remove the limitation. Until then, prefer `"projection"` at that scale,
-where the stall does not bite (below).
+`mass="free"`. After every step the log-weights are projected exactly onto
+the capped simplex `{w : sum(w) = D, w <= cap * d}`, and the closing float64
+projection that both parametrizations share makes the returned vector exact.
+It requires `mass="conserve"` and no L0 gates.
+
+**The per-step cap projection.** It is the Kullback-Leibler projection: one
+common shift `s` for every log-weight, then a clamp to each log cap,
+`log_w <- min(log_w + s, log(cap * d))`. Uncapped records keep their
+relative weights, the direction the softmax cannot see, and only records
+whose share would pass their caps are held at them. `s` solves
+`sum(min(exp(log_w + s), cap * d)) = D`, whose left side is continuous and
+nondecreasing in `s`. Active-set rounds find it: with `C` the records already
+at their caps, the next shift solves
+`sum(cap * d[C]) + exp(s) * sum(exp(log_w)[not C]) = D`. 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. If eight rounds pass without that, a sort over the
+remaining records finishes it, which bounds the worst case at `O(n log n)`.
+Both paths compute the shift in float64. The same bound keeps
+`D - sum(cap * d[C])` at least `exp(s) * sum(exp(log_w)[not C])`, the
+uncapped records' share, so when rounding takes it to zero or below, the
+capped records' caps already hold `D`: the rounds stop at the current shift,
+and the other records keep a share below float64 rounding of `D`. When the
+uncapped weights sum below float64's normal range, the rounds hand over to
+the sort, which sums them in log space, so a weight that underflows float64
+still counts. The solve records
+`options["iterate_selection_receipt"]["softmax_in_loop_max_cap_ratio"]`, 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.
+`test_softmax_cap_projection.py` pins the projection for every drawn input,
+including zero weights and zero caps drawn independently, a cap of exactly
+1, and inputs where some records' caps meet the total while the rest sit
+below float64 rounding of it or below float64's range. It conserves the
+total, respects every cap, is idempotent and invariant to a constant added
+to the log-weights, and moves an input already within its caps by the shift
+alone. It also meets the Kullback-Leibler projection's KKT conditions,
+equals a water-fill written from the definition in 60-digit decimal
+arithmetic, and gives the same answer by its active-set and sorted paths,
+including on an input built to need one round per record. Wherever the cap
+step that the full-scale measurement below ran settles in its active-set
+rounds, with every nonzero weight in float64's normal range, the test checks
+that the current step returns the same shift and round count, bit for bit.
+
+**What the projection replaced.** Before, each step renormalized and clamped
+in a loop of at most 32 rounds. On the ACS local release the rounds ran out on
+most epochs, and how far that left the in-loop weights past their caps was
+not recorded. It has since been measured
+(`experiments/us-acs-local-l2-basis-20260928/results/exact_cap.md` and
+`results/exact_cap_float32.json`):
+
+- **Why the loop ran out.** Each round's shift was `log(D)` in float64 minus a
+  float32 `logsumexp`. On a vector already within its caps that difference
+  does not reach zero. It bottoms out at the float32 rounding residual of
+  `log(D)` itself, 7.8e-8 here. That shift lifts the one to three capped
+  records with log caps below 2 in magnitude one float32 step over their caps,
+  the clamp puts them back, and the next round repeats it. On full-scale vectors
+  the projection had already made exact, the loop still ran out in 10 of 12
+  trials.
+- **How far past the cap it was.** Not materially. Two full-scale
+  configurations whose earlier runs ran out were rerun under both steps: share
+  0.5 at `l2_lambda = 0.03`, and share 0.9 unpenalized. The largest in-loop
+  weight over its cap was 1 + 1.2e-5 and 1 + 8.7e-6 under the loop, and
+  1 + 9.0e-6 under the projection at both. The loop ran out on 713 and 705 of
+  800 steps. The projection settled in one to three rounds on every step and
+  took 12.7-22.5 ms a step, against the loop's 48.7-53.8 ms.
+- **What sets that residual.** The float32 softmax. The projected
+  log-weights sit within their float32 log caps, but the solve realizes
+  `D * softmax(log_w)` with a float32 normalizer over 1.59M records, whose
+  total was off by up to 8.2e-6. The cap ratio exceeded that total error by
+  at most 6.1e-7. With a float64 normalizer on the same log-weights the largest
+  ratio was 1 + 4.9e-7, the float32 rounding of the log-weights themselves.
+  That decomposition covers only the projected vectors; the same softmax
+  plausibly sets the loop's residual, but no run measured it.
+- **The solves.** National ESS agreed within 0.1% between the two steps
+  (21,834 against 21,835, and 93,609 against 93,598). At share 0.5 and
+  `l2_lambda = 0.03` the projection's final loss was 0.01899. Three loop runs
+  of that configuration gave 0.01910-0.01917, so the projection's loss is
+  0.6-0.9% lower, from one run. At share 0.9 the loss was 0.2% lower (0.03129
+  against 0.03137).
 
 Evidence:
 
@@ -165,8 +220,10 @@ recalibrated from its own checkpoint (`experiments/us-acs-local-l2-basis-2026092
   `l2_lambda` is not an equal-ESS comparison. At `0.03` the held-out
   comparison is a toss-up: softmax has slightly lower capped error, projection
   more held-out targets within 10% on both folds. Projection's training loss is
-  6% lower there, against 0.3% run-to-run variation. With that and the
-  cap-loop limitation above, projection is the recommended parametrization.
+  6% lower there, against 0.3% run-to-run variation, so projection is the
+  recommended parametrization. Rerun with the exact cap projection, softmax's
+  training loss at `0.03` is 0.01899 (one run), still 5% above projection's
+  0.01806.
 - **Record basis.** The record basis at `0.1` lowers national ESS from 13,646
   to 8,689, as its `w ∝ d ** 2` optimum predicts.
 - **Chi-square basis.** The chi-square basis raises ESS smoothly with
@@ -191,7 +248,7 @@ result = calibrate(
     l2_lambda=...,
     l2_basis="chi_square",
     # mass_parametrization="softmax" for problems where every target presses
-    # the same way; see the cap-loop limitation above.
+    # the same way (above).
 )
 result.options["l2_basis"], result.options["mass_parametrization"]
 result.chi_square_distance, result.effective_sample_size
@@ -200,8 +257,9 @@ result.chi_square_distance, result.effective_sample_size
 The ACS local-area tool takes the same settings as
 `--l2-lambda`, `--l2-basis {record,chi_square}` and
 `--mass-parametrization {projection,softmax}`. It records them, with the
-realized chi-square distance, in `calibration_summary.json` and in the build
-manifest's `calibration` block. Both join the tool's solver-settings stamp,
+realized chi-square distance and, under softmax, the last epoch batch's
+largest in-loop weight over its cap, in `calibration_summary.json` and in the
+build manifest's `calibration` block. Both join the tool's solver-settings stamp,
 so `--resume` and the already-complete shortcut refuse weights solved under
 any other solver setting (cap, loss cap, `l2_lambda`, `l2_basis`,
 `mass_parametrization`, seed, epoch batch); a stamp written before these two

docs/evidence/spec-engine/us-f0-coverage.json

@@ -1642,13 +1642,13 @@
           "compiler_ir.node_slices"
         ],
         "expected": {
-          "map_sha256": "07fe9086b57a15ab1c2551fce5c30aebfa55175d76a308f7dbb268ac49da747a",
-          "protocol_sha256": "4ff992324004d4c09d4165bc8d24ccff15b9d84858836851b12d532c870920f4"
+          "map_sha256": "7ef47df73404e40d20e12292c35bc81ce18620e4608f8a52ce2b8f95ec32f4d1",
+          "protocol_sha256": "76c372eefaf7581557d530ec68123802b974860d9b2e581f45b30aa55cedfbf5"
         },
         "failures": [],
         "observed": {
-          "map_sha256": "07fe9086b57a15ab1c2551fce5c30aebfa55175d76a308f7dbb268ac49da747a",
-          "protocol_sha256": "4ff992324004d4c09d4165bc8d24ccff15b9d84858836851b12d532c870920f4"
+          "map_sha256": "7ef47df73404e40d20e12292c35bc81ce18620e4608f8a52ce2b8f95ec32f4d1",
+          "protocol_sha256": "76c372eefaf7581557d530ec68123802b974860d9b2e581f45b30aa55cedfbf5"
         },
         "status": "covered"
       },
@@ -1663,7 +1663,7 @@
           "compiler_ir.seed_stream_map"
         ],
         "expected": {
-          "implementation_sha256": "4ff992324004d4c09d4165bc8d24ccff15b9d84858836851b12d532c870920f4",
+          "implementation_sha256": "76c372eefaf7581557d530ec68123802b974860d9b2e581f45b30aa55cedfbf5",
           "protocol": "legacy-v1",
           "streams": [
             "build_model",
@@ -1684,7 +1684,7 @@
         },
         "failures": [],
         "observed": {
-          "implementation_sha256": "4ff992324004d4c09d4165bc8d24ccff15b9d84858836851b12d532c870920f4",
+          "implementation_sha256": "76c372eefaf7581557d530ec68123802b974860d9b2e581f45b30aa55cedfbf5",
           "protocol": "legacy-v1",
           "streams": [
             "build_model",
@@ -2585,7 +2585,7 @@
       "country": "us",
       "schema_id": "country_spec",
       "schema_version": 1,
-      "spec_sha256": "0c713aabfd0afc3703415fea86a8210a3f4346918f774f371fb16aa663348c1d"
+      "spec_sha256": "874294b24bf70671b3ec78bb1d1a2c3461fd585e2af1388265c27be26c18d0ba"
     }
   },
   "report_schema_version": 3,
@@ -2595,7 +2595,7 @@
     "country": "us",
     "schema_id": "country_spec",
     "schema_version": 1,
-    "spec_sha256": "0c713aabfd0afc3703415fea86a8210a3f4346918f774f371fb16aa663348c1d"
+    "spec_sha256": "874294b24bf70671b3ec78bb1d1a2c3461fd585e2af1388265c27be26c18d0ba"
   },
   "status": "pass"
 }