Skip to content

docs: document that l0_coefficient saturates, and warn when it has - #720

Open
guptaishaan wants to merge 6 commits into
decoderesearch:mainfrom
guptaishaan:docs/jumprelu-l0-coefficient-saturation
Open

docs: document that l0_coefficient saturates, and warn when it has#720
guptaishaan wants to merge 6 commits into
decoderesearch:mainfrom
guptaishaan:docs/jumprelu-l0-coefficient-saturation

Conversation

@guptaishaan

Copy link
Copy Markdown

Addresses #719. Builds on #718 — that PR frees the JumpReLU threshold to move; this is about what happens once it can.

Note on the diff: this branch is stacked on #718, so it carries that PR's four commits as well. Only the last commit (docs: document that l0_coefficient saturates, and warn when it has) is new here. It needs #718 because on main the threshold is a read-only exp(log_threshold) property that barely moves, so the diagnostic below can't apply. Happy to rebase once #718 lands, or to retarget it if #718 changes shape.

#719 reports that l0_coefficient does nothing across 1/10/100. It isn't inert, it saturates. Sweeping downward, on cached activations with lr=3e-4, bandwidth 0.05, init threshold 0.01, 2000 steps:

l0_coefficient 0 1e-3 1e-2 0.1 1 10 100
final L0 478 364 124 36 26.6 26.4 26.0

The knob moves L0 10x over [1e-3, 1e-1] and is flat above ~1, so the values in the report were all past the point where it stops doing anything.

The threshold receives c * g_l0 + g_mse, and Adam fixes the step size at ~lr, so c only sets the direction of the update. Measured mid-training, the unit-c l0 gradient on the threshold is ~1000x the reconstruction gradient, so by c ≈ 0.01 the direction is already entirely the l0 term and further increases are no-ops.

Where that boundary sits depends on activation scale, jumprelu_bandwidth, d_sae, batch size and n_steps — at 20k steps instead of 2k the same sweep is still responsive at c=1 → 10. So a better default wouldn't fix it.

This PR doesn't change any optimization behaviour. It:

  • documents the saturation, scoped to step mode — in tanh mode the penalty also back-propagates through W_enc/b_enc/W_dec and l0_coefficient responds normally (L0 3.84 → 2.59 → 0.84 across c = 1/10/100), so the "only gradient path is the threshold" reasoning doesn't apply there
  • warns at the end of training when the threshold travelled ≥90% of what Adam permits, which is the signal that the run is in the saturated regime

On the sweep above the warning fires at c = 1 and 100 and stays quiet at 0.01 and 0.1, which is the intended split.

Implementation is a no-op TrainingSAE.training_convergence_warning() hook, so other architectures are unaffected. It's skipped when resuming from a checkpoint, since travel is measured from initialization.

Two things I deliberately left alone: the l0_coefficient=1.0 default (a behaviour change, and per the above no fixed value is right across setups), and the 5.0 in the tanh example, which isn't affected.

Happy to take this further — a target-L0 scheme is the option that would actually be robust to the scale-dependence — but that's a much larger change and a config-semantics decision, so I've kept this to documenting what the knob currently does.

Tests: 8 new, CPU-only, no model download. Existing suite passes.

Hetens and others added 6 commits July 20, 2026 16:15
Under Adam the per-step parameter movement is capped at ~lr regardless of
gradient magnitude, so exp(log_threshold) moves the effective threshold by
only ~threshold * lr per step. With the default 0.01 init the threshold is
effectively frozen and the L0 penalty cannot reduce L0 (decoderesearch#494). Reference
implementations (DeepMind pseudocode, dictionary_learning) use a direct
threshold parameter; dictionary_learning explicitly notes poor results
with log_threshold.

Old checkpoints saving log_threshold are converted on load.
The threshold regression test asserts it travels at the Adam step size
rather than being scaled down by its own magnitude; it fails on the old
log_threshold parameterization (0.0003 moved vs 0.01 required).
l0_coefficient is documented as the knob for the JumpReLU
reconstruction/sparsity tradeoff, but past a certain value it stops
having any effect (decoderesearch#719).

The threshold receives c * g_l0 + g_mse, and Adam fixes the step size at
~lr, so c only sets the direction of the update. Once the l0 term
dominates that direction, raising c further changes nothing. The point
where that happens depends on activation scale, bandwidth, d_sae, batch
size and n_steps, so it can't be fixed by choosing a better default.

This doesn't change any optimization behaviour. It documents the
saturation for "step" mode ("tanh" mode back-propagates through the
encoder and decoder and responds to c normally), and warns at the end of
training when the threshold travelled ~as far as Adam allows, which is
the signal that the run is in the saturated regime.

Implemented as a no-op TrainingSAE.training_convergence_warning() hook so
other architectures are unaffected.
Conflict in sae_trainer.py fit(): decoderesearch#682 reformatted the tqdm call to pass
initial=self.n_training_samples, adjacent to the started_from_scratch
guard added here. Kept both.

decoderesearch#682 also makes resume_from_checkpoint actually take effect, so
n_training_steps is now genuinely non-zero on a resumed run and the guard
does the job it was written for.
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.

2 participants