feat(adamw): let a recipe exempt parameters from decoupled weight decay - #2030
feat(adamw): let a recipe exempt parameters from decoupled weight decay#2030ooples wants to merge 1 commit into
Conversation
WeightDecay is a single coefficient applied to the whole flat parameter vector, so there was no way to express "decay these parameters but not those". Published recipes routinely need exactly that. RecurrentGemma (Botev et al., 2024) Section 2: "we do not apply weight decay to the parameters of the recurrent (RG-LRU) layers during training." The ordinary transformer recipe exempts biases and normalization gains for the same reason -- decay pulls weights toward zero, which helps where magnitude is the thing being regularized and corrupts parameters whose VALUE carries meaning, such as a recurrence's own decay rate. WeightDecayMask is an optional per-parameter multiplier on the decay term. Null decays everything, which is exactly what every caller had before, so no existing model changes behaviour. Entries multiply elementwise: 1 decays normally, 0 exempts, fractions scale. The gradient update is untouched; this is decoupled decay only. Applied at all three eager sites -- the two vector paths and the element-wise span loop -- so they cannot disagree with each other. TryGetFusedOptimizerConfig now DECLINES when a mask is set. The fused config carries decay as a single float, so a masked run cannot be expressed in it, and the compiled kernel would go on decaying the parameters the eager path exempts. Declining is the same mechanism that method already uses for adaptive learning rates. Silently diverging between the two paths would be the worst of the available options, since it would only show up as a slow accuracy drift. Four tests: the null default decays everything, zero entries exempt exactly those parameters, fractional entries scale the decay proportionally, and a masked optimizer reports no fused config while an unmasked one still fuses. Each isolates the decay term by stepping with a ZERO gradient, so anything that moves a parameter is decay and only decay. 75 optimizer tests pass unchanged, including FusedSpecMatchesEagerBehaviour and the copy-constructor suite. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Warning Review limit reached
Next review available in: 26 minutes Limit details: You’ve used the included review currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?Wait for the limit to reset, then comment An organization admin can change what happens after included review limits in Billing. How do review limits work?CodeRabbit enforces per-developer PR review limits within each organization. For paid Pro and Pro+ reviews, CodeRabbit uses a developer's included PR review attempts over the past 7 days to set the current hourly allowance. At typical activity levels, the full plan allowance applies. Higher sustained activity can lower the allowance until earlier attempts leave the 7-day window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Superseded by #2029. This PR described itself as a follow-up to #2029, so keeping it on a separate branch just split one review across two places. The Closing here; review continues on #2029. |
1 participant
Follow-up to #2029, which needed this and could not have it.
Why
WeightDecayis a single coefficient applied to the whole flat parameter vector, so there is no way to express "decay these parameters but not those". Published recipes routinely need exactly that:The reason is the same in both cases: decay pulls weights toward zero, which helps where magnitude is what's being regularized, and corrupts parameters whose value carries meaning — a recurrence's own decay rate, for instance.
In #2029 I left this measure out rather than ship a property that silently did nothing. This is the capability that makes it expressible.
What changed
WeightDecayMask— an optional per-parameter multiplier on the decay term.nulldecays everything, which is exactly what every caller had before. No existing model changes behaviour.1decays normally,0exempts, fractions scale.The fused path
TryGetFusedOptimizerConfignow declines when a mask is set.The fused config carries decay as a single
float, so a masked run cannot be represented in it — the compiled kernel would keep decaying the parameters the eager path exempts. Declining is the same mechanism that method already uses for adaptive learning rates.Silently diverging between the compiled and eager paths would be the worst option available: it would surface only as a slow accuracy drift, with nothing pointing at the optimizer.
Verification
NullMask_DecaysEveryParameter...ZeroMaskEntries_ExemptExactlyThoseParametersFractionalMaskEntries_ScaleTheDecayMaskPresent_DeclinesFusedCompilationEach isolates the decay term by stepping with a zero gradient, so anything that moves a parameter is decay and only decay.
75 optimizer tests pass unchanged, including
FusedSpecMatchesEagerBehaviourand the copy-constructor suite — this touches an optimizer every model uses, so that regression check is the point.Note for reviewers
This is shared-optimizer code. The two things that make it safe are the
nulldefault (byte-identical behaviour for every current caller) and the fusion opt-out, and both are directly tested rather than asserted.A natural next step, not included here: have
RecurrentGemmaLanguageModelbuild a mask that zeroes itsRealGatedLinearRecurrenceLayerslots, which would complete the paper's third measure. That needs a layer-to-flat-slot mapping and is worth its own PR.🤖 Generated with Claude Code