Skip to content

Fix AdaSmooth NaN on zero-initialized parameters - #522

Merged
kozistr merged 1 commit into
kozistr:mainfrom
winklemad:fix-adasmooth-zero-init-nan
Sep 9, 2026
Merged

kozistr merged 1 commit into
kozistr:mainfrom
winklemad:fix-adasmooth-zero-init-nan

Conversation

@winklemad

Copy link
Copy Markdown
Contributor

Problem

AdaSmooth computes its efficiency ratio by dividing the running signed movement by the running absolute movement:

c = s.sum().abs_().div_(n.sum())  # e_t

On the first step, a zero-initialized parameter (a bias, or any zero-init layer — very common) has not moved yet: p_diff is zero, so n.sum() is 0 and the update is 0 / 0 = NaN, which permanently corrupts the parameter.

import torch
from pytorch_optimizer import AdaSmooth
p = torch.nn.Parameter(torch.zeros(4))
opt = AdaSmooth([p], lr=1e-2)
opt.zero_grad(); ((p - 1.0) ** 2).sum().backward(); opt.step()
print(p)  # tensor([nan, nan, nan, nan]) on main

The convergence tests only use randomly-initialized parameters, so this was never hit.

Fix

Guard the denominator with eps (n.sum() + eps) so a zero total movement yields a finite ratio; the first-step update is then well-defined and the parameter trains normally. For non-zero parameters eps is negligible, so behavior is unchanged.

Added test_adasmooth_zero_initialized_parameter (NaN on the current code, finite after the fix). ruff clean; the existing adasmooth tests pass.

AdaSmooth's efficiency ratio divides by the running sum of absolute
parameter movement, `c = |s.sum()| / n.sum()`. On the first step a
zero-initialized parameter (a bias, or any zero-init layer) has not
moved, so `p_diff` is zero, `n.sum()` is zero, and the update becomes
0 / 0 = NaN, permanently corrupting the parameter. Since the convergence
tests use randomly-initialized parameters, this was never exercised.

Guard the division with eps so a zero movement yields a finite ratio.
@winklemad
winklemad requested a review from kozistr as a code owner September 9, 2026 06:30
@github-actions github-actions Bot added the optimizer about optimizer label Sep 9, 2026
@codecov

codecov Bot commented Sep 9, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (d880323) to head (70c0c41).

Additional details and impacted files
@@            Coverage Diff            @@
##              main      #522   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files          128       128           
  Lines        12444     12444           
=========================================
  Hits         12444     12444           
Flag Coverage Δ
3.10 100.00% <100.00%> (ø)
3.11 100.00% <100.00%> (ø)
3.12 100.00% <100.00%> (ø)
3.13 100.00% <100.00%> (ø)
3.14 100.00% <100.00%> (ø)
3.8 99.98% <100.00%> (ø)
3.9 99.98% <100.00%> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@kozistr kozistr left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for your contribution! looks good to me!

@kozistr
kozistr merged commit 9f74878 into kozistr:main Sep 9, 2026
17 checks passed
@winklemad

Copy link
Copy Markdown
Contributor Author

Thanks @kozistr!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

optimizer about optimizer size/XS

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants