#1034·ai-toolkit

EMA shadow weights diverge — `sub_` changed to `add_` in commit 4a41f14

Author: kekm8-cpuCreated Sep 10, 2026Updated Sep 12, 2026

Commit 4a41f14 ("Fix feedback when doing an ema with feedback so it uses its own scale") flipped the sign of the EMA shadow update in ExponentialMovingAverage.update.

Before:

python
tmp = (s_param_float - param_float)
tmp.mul_(one_minus_decay)
s_param_float.sub_(tmp)

After:

python
gap = (s_param_float - param_float)
s_param_float.add_(gap * one_minus_decay)

Since gap = s - p, the correct update is s -= (1 - decay) * gap. With add_, the shadow becomes:

s_new = s + (1 - decay) * (s - p) = 1.01 * s - 0.01 * p   (at decay=0.99)

so the shadow is pushed away from the parameters instead of toward them, multiplying by 1.01 every step.

Impact: any run with use_ema: true saves diverging shadow weights. Training itself is unaffected — the live parameters and the loss curve look completely normal — so there is no visible signal that anything is wrong. At decay=0.99 the growth is ~12x per 250 steps, which exceeds bf16's usable range within a few hundred steps and produces LoRAs that render as noise at any strength multiplier.

Observed, LoRA training on Krea 2, decay 0.99, bf16 save, max |weight| per checkpoint:

step max abs weight ratio vs previous
1500 1.082e4
1750 1.306e5 12.07
2000 1.573e6 12.04
2019 1.876e6 1.19 (19 steps)

1.01^250 = 12.03 and 1.01^19 = 1.21, matching the observed ratios. An identical config run on b1bf3e4 (before this commit) stays at ~0.04 as expected.

Fix — one line:

python
gap = (s_param_float - param_float)
s_param_float.sub_(gap * one_minus_decay)

The use_feedback branch below it is correct as written; param_float.add_(gap * self.feedback_rate) pulls the parameter toward the shadow, which is the intended direction.