timm's AdafactorBigVision optionally clips each update by its root-mean-square: the update should
be divided by max(1, RMS/threshold), a contraction that shrinks an over-sized update to the
threshold and leaves a small one alone. The code inverts this two ways at once. It divides by
RMS * threshold instead of RMS / threshold, and it clamps the divisor with max=1.0 instead of
min=1.0, so the divisor is always at most one. The clip therefore never shrinks a large update and
instead amplifies every small update up to RMS = 1/threshold. A gradient spike, exactly what the
clip exists to tame, passes through unbounded. The correct sibling adafactor.py implements the
same clip properly, so this is a divergence bug, not a deliberate choice.
timm, 1.0.28 and current main (byte-identical):
# timm/optim/adafactor_bv.py, lines 273-276
# Clip by RMS value
if clipping_threshold is not None:
denom = (update.norm(2) / ((update.numel() ** 0.5) / clipping_threshold)).clamp_(max=1.0)
update.div_(denom)
# timm/optim/adafactor.py, line 208 (the correct sibling)
update.div_((self._rms(update) / group['clip_threshold']).clamp_(min=1.0))update.norm(2) / (sqrt(N) / c) is RMS(update) * c, and clamp_(max=1.0) makes the divisor
min(1, RMS*c) <= 1, so update.div_ never shrinks and can amplify. The sibling divides by
max(1, RMS/c) >= 1, the true contraction. The big_vision line has both the reciprocal and the
clamp direction flipped relative to its own sibling.
Driving the real clip line against the real sibling on a synthetic update, and the whole optimizer under a gradient spike (fp64):
exact clip line on a synthetic update, threshold c=1.0 (result RMS must be <= 1):
update RMS 0.3: adafactor_bv -> 1.000 sibling adafactor -> 0.300
update RMS 3.0: adafactor_bv -> 3.000 sibling adafactor -> 1.000
update RMS 5.0: adafactor_bv -> 5.000 sibling adafactor -> 1.000
real AdafactorBigVision, gradient spike after warmup (applied update RMS):
clip disabled (None) : 2.048
clip enabled (1.0) : 2.048 <- unchanged, clip did nothing
A small update at RMS 0.3 is amplified to 1.0 where it should be left alone, and a large update
at RMS 3 or 5 is passed through unchanged where it should be clipped to 1. End to end, after
warming the second-moment estimate on tiny gradients, a gradient spike produces an update of
RMS 2.048, and enabling the clip leaves it at 2.048, identical to no clip: the clip is a no-op
exactly when it should protect training. The correct sibling clips the same updates to the threshold.
The clip runs when clipping_threshold is set. In timm it defaults to None, so it is off unless
enabled, but AdafactorBigVision exists to reproduce big_vision "Scaling Vision Transformers"
training, whose scale_by_adafactor ships clipping_threshold = 1.0 on by default, so enabling it
is the intended recipe rather than a misconfiguration; clipping_threshold is a documented
first-class constructor argument and 1.0 is an in-contract value. When enabled the failure is
silent and severe: the RMS clip, a stability safeguard, is inverted, so gradient spikes are never
tamed and small updates are inflated. The fix is to divide by max(1, RMS/threshold) as the sibling
does.
excerpt.py: the big_vision clip and the correct sibling clip quoted with the flags that name the fault, and theclipping_thresholddefault and big_vision context (Apache-2.0).clip.py: the correct and buggy divisors modelled on the RMS, so the correct clip is a contraction bounding the result by the threshold and the buggy one is an expansion.consequence.py: the real clip line versus the real sibling on synthetic updates, and the whole optimizer under a gradient spike where enabling the clip changes nothing.test_clipflip.py: the model's correct clip contracts while the buggy one expands, the divisors are the reciprocal-and-clamp-flipped forms, the real clip line fails where the sibling succeeds, and the real gradient-spike update is unchanged by enabling the clip.
python clip.py
python consequence.py
python test_clipflip.py
The clip lines are quoted from current main; the numbers are produced by the real optimizer. The
fix is to divide by max(1, RMS/threshold), so the clip bounds the update RMS by the threshold as
the sibling and the reference big_vision implementation do.