Skip to content

v-code01/clipflip

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

6 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

clipflip

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.

The code

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.

What it does

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.

Scope

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.

Layout

  • excerpt.py: the big_vision clip and the correct sibling clip quoted with the flags that name the fault, and the clipping_threshold default 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.

Reproduce

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.

About

timm AdafactorBigVision RMS update clip is inverted: divides by min(1, RMS*threshold) instead of max(1, RMS/threshold), so it never clips a large update and amplifies small ones; sibling adafactor.py clips correctly

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages