Skip to content

feat(adversary): stateless sign-ascent SRC_STEP for persistent-PGD sources (SPEC §6)#943

Open
ocg-goodfire wants to merge 2 commits into
feature/jaxfrom
bridge/task-lp-sign-ascent-source
Open

feat(adversary): stateless sign-ascent SRC_STEP for persistent-PGD sources (SPEC §6)#943
ocg-goodfire wants to merge 2 commits into
feature/jaxfrom
bridge/task-lp-sign-ascent-source

Conversation

@ocg-goodfire

Copy link
Copy Markdown
Collaborator

What

Implements the sign SRC_STEP variant already named as a valid instantiation in SPEC §6: sources += lr·sign(grad) with the same S15 clamp projection, selected via optimizer: {type: sign, lr_schedule: ...} (a discriminated union with adam on PersistentPGDReconLossConfig.optimizer).

The variant is stateless: PersistentAdversary.opt_state is now SourcesAdamState | None (None iff sign), so the source Adam m/v buffers never exist. On full32L that is −19 GiB/GPU resident (2× the 9.5 GiB source bundle); combined with the validated source_dtype: bfloat16 (#938) the source bundle goes 28.5 → ~4.75 GiB.

Invariants (SPEC IDs)

  • §6 SRC_STEP: sign was already a named valid instantiation — no SPEC amendment. Production stays ★ adam.
  • S13'/S14': lifecycle unchanged — n_warmup supplemental ascents + one final ascent from the main backward, coeff-unscaled; sources_ascend_project dispatches on the optimizer config.
  • S15: same clamp-to-[0,1] after every ascent.
  • N1 (fp32 SRC_STEP moments): vacuous for sign — no moments. The adam path (sources_adam_ascend_project) is bit-untouched.

Configs

Four Tier-1 quality A/B arms (pile-4L PPGD-bsc, 10k steps, same seed, dp 8): adam control + sign at lr {0.003, 0.01, 0.03} — a step-size sweep since Adam's normalized per-update magnitude ≈ lr was tuned at 0.01, but sign steps are exactly lr. All arms fp32 source to isolate the update rule. Results will be logged on the bridge task (lp-sign-ascent-source).

Validation

  • pytest param_decomp/tests/ green at default device count; adversary/checkpoint tests also green at XLA_FLAGS=--xla_force_host_platform_device_count=4.
  • New: test_sign_src_step.py (unit: ±lr move, clamp, dtype, statelessness) + test_step_trains_with_sign_src_step (full jitted train step: losses finite, sources move and stay projected, opt_state stays None).
  • make check clean (ruff + basedpyright 0 errors).

🤖 Generated with Claude Code

https://claude.ai/code/session_01ALaUCqhksKmWNbxjfmuMUi

ocg-goodfire and others added 2 commits July 3, 2026 14:57
…urces (SPEC §6)

Implements the `sign` SRC_STEP variant already named in SPEC §6:
`sources += lr·sign(grad)` with the same S15 clamp projection, configured as
`optimizer: {type: sign, lr_schedule: ...}` (discriminated union with `adam`
on PersistentPGDReconLossConfig.optimizer).

The variant is stateless: PersistentAdversary.opt_state is SourcesAdamState |
None (None iff sign), so the Adam m/v buffers never exist — on full32L that
is −19 GiB/GPU resident (2× the 9.5 GiB source bundle), on top of the
validated bf16-source win (source bundle 28.5 → ~4.75 GiB combined).

S13'/S14' lifecycle (n_warmup supplemental ascents + one final ascent from
the main backward, coeff-unscaled) is unchanged; sources_ascend_project
dispatches on the optimizer config and the adam path
(sources_adam_ascend_project) is bit-untouched. Production remains ★ adam.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ALaUCqhksKmWNbxjfmuMUi
…4L PPGD-bsc, 10k steps)

Four arms off pile_ppgd_bsc.yaml (steps 10k, dp 8, slow_every 2000, same
seed): adam control (production ★) + sign at lr {0.003, 0.01, 0.03} — a
step-size sweep since Adam's normalized update ≈ lr was tuned at 0.01 but
sign steps are exactly lr. All arms fp32 source to isolate the update rule.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01ALaUCqhksKmWNbxjfmuMUi
@ocg-goodfire

Copy link
Copy Markdown
Collaborator Author

Tier-1 quality A/B result: sign-ascent FAILS the gate at all three step sizes (pile-4L PPGD-bsc, 10k steps, same seed, fp32 sources; runs p-e6f2e806 adam / p-992df179 lr.003 / p-41cee32f lr.01 / p-07b7e68b lr.03).

The adversary-independent robustness probe (fresh-PGD 20-step eval) vs adam control:

  • lr 0.003: +400–900% throughout — sign adversary far too weak, decomposition never faces real pressure (its better masked CE/KL is the weak-adversary confound, not a win)
  • lr 0.01: metastable — oscillates between parity and +300–620% blowups, widening over training (adversary falls behind the co-evolving decomposition)
  • lr 0.03: +15–27% consistently, plus +10–25% masked CE/KL and +15–30% CI-L0 — dominated on every axis

Conclusion: Adam's adaptive per-coordinate scaling is load-bearing for the persistent adversary (3 ascents/step tracking a moving target — unlike fresh PGD's 20 dedicated sign steps from scratch). The −19 GiB m/v saving is not available via naive sign; bf16 source+moments (#938, −14.25 GiB) remains the validated path.

Full writeup: lore 2026-07-03--sign-ascent-ppgd-source-fails-quality-gate.

Merge-or-close is the team's call: the impl is a faithful SPEC §6 sign instantiation with tests, but there is now no production use for it. Closing is consistent with the no-dead-code norm; merging keeps the variation point exercised for future signum-style experiments.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant