docs: document that l0_coefficient saturates, and warn when it has - #720
Open
guptaishaan wants to merge 6 commits into
Open
docs: document that l0_coefficient saturates, and warn when it has#720guptaishaan wants to merge 6 commits into
guptaishaan wants to merge 6 commits into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Addresses #719. Builds on #718 — that PR frees the JumpReLU threshold to move; this is about what happens once it can.
#719 reports that
l0_coefficientdoes nothing across 1/10/100. It isn't inert, it saturates. Sweeping downward, on cached activations withlr=3e-4, bandwidth 0.05, init threshold 0.01, 2000 steps:l0_coefficientThe 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, soconly sets the direction of the update. Measured mid-training, the unit-cl0 gradient on the threshold is ~1000x the reconstruction gradient, so byc ≈ 0.01the 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 andn_steps— at 20k steps instead of 2k the same sweep is still responsive atc=1 → 10. So a better default wouldn't fix it.This PR doesn't change any optimization behaviour. It:
stepmode — intanhmode the penalty also back-propagates throughW_enc/b_enc/W_decandl0_coefficientresponds 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 thereOn the sweep above the warning fires at
c = 1and100and stays quiet at0.01and0.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.0default (a behaviour change, and per the above no fixed value is right across setups), and the5.0in thetanhexample, 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.