From 6654c58dabb8af4b48fda11b1006ad1f5e698428 Mon Sep 17 00:00:00 2001 From: shanyu Date: Sat, 26 Sep 2026 00:29:10 +0800 Subject: [PATCH] Apply LAMB debias to the moments inside the trust ratio debias=True multiplied the learning rate by the Adam correction and left adam_norm on the raw moments. Weight decay, which is added before the norm, was scaled by that factor too, so the trust ratio did not match a debiased step (issue 443). The correction is now m / (1 - beta1**t) and v / (1 - beta2**t) on copies of the moments. The stored averages are unchanged, and the default debias=False path is the paper's uncorrected step. --- tests/test_lamb_debias.py | 39 +++++++++++++++++++++++++++++++++++++++ torch_optimizer/lamb.py | 20 +++++++++++--------- 2 files changed, 50 insertions(+), 9 deletions(-) create mode 100644 tests/test_lamb_debias.py diff --git a/tests/test_lamb_debias.py b/tests/test_lamb_debias.py new file mode 100644 index 00000000..8115b60c --- /dev/null +++ b/tests/test_lamb_debias.py @@ -0,0 +1,39 @@ +"""debias=True corrects the moments inside the trust ratio.""" +import torch +import torch_optimizer as optim + + +def test_default_debias_is_the_uncorrected_step(): + param = torch.nn.Parameter(torch.tensor([2.0, -1.0])) + opt = optim.Lamb([param], lr=0.1, weight_decay=0.0, betas=(0.9, 0.999)) + grad = torch.tensor([1.0, 0.5]) + param.grad = grad.clone() + opt.step() + # m = (1-0.9) g, v = (1-0.999) g^2, step = m / sqrt(v) + direction = (0.1 * grad) / (0.001 * grad * grad).sqrt() + weight_norm = torch.tensor([2.0, -1.0]).norm() + trust = weight_norm / direction.norm() + expected = torch.tensor([2.0, -1.0]) - 0.1 * trust * direction + assert torch.allclose(param.detach(), expected, atol=1e-5) + + +def test_debias_corrects_adam_norm(): + # One step, scalar parameter. Stored moments must stay the raw EMA. + param = torch.nn.Parameter(torch.tensor([2.0])) + opt = optim.Lamb( + [param], lr=0.1, weight_decay=0.1, betas=(0.9, 0.999), debias=True + ) + param.grad = torch.tensor([1.0]) + opt.step() + state = opt.state[param] + assert torch.allclose(state["exp_avg"], torch.tensor([0.1])) + assert torch.allclose(state["exp_avg_sq"], torch.tensor([0.001])) + m_hat = 0.1 / (1 - 0.9) + v_hat = 0.001 / (1 - 0.999) + direction = m_hat / (v_hat ** 0.5) + 0.1 * 2.0 + # scalar norms are absolute values; trust ratio is |p| / |direction| + updated = 2.0 - 0.1 * (2.0 / abs(direction)) * direction + assert torch.allclose(param.detach(), torch.tensor([updated]), atol=1e-5) + # Scaling the learning rate and leaving adam_norm uncorrected gives + # about 1.937 on this step, not the corrected 1.8. + assert abs(param.detach().item() - 1.936754) > 0.05 diff --git a/torch_optimizer/lamb.py b/torch_optimizer/lamb.py index 7e948b05..3e5b2d2f 100644 --- a/torch_optimizer/lamb.py +++ b/torch_optimizer/lamb.py @@ -1,5 +1,3 @@ -import math - import torch from torch.optim.optimizer import Optimizer @@ -126,19 +124,23 @@ def step(self, closure: OptLossClosure = None) -> OptFloat: # v_t exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) - # Paper v3 does not use debiasing. + # Paper v3 does not use debiasing. When it is requested, Adam's + # correction is applied to the moments that enter the trust + # ratio. Scaling the learning rate instead leaves adam_norm + # uncorrected and also scales the weight-decay term (issue 443). + # The stored moments are not divided in place. if self.debias: - bias_correction = math.sqrt(1 - beta2 ** state["step"]) - bias_correction /= 1 - beta1 ** state["step"] + exp_avg_hat = exp_avg / (1 - beta1 ** state["step"]) + exp_avg_sq_hat = exp_avg_sq / (1 - beta2 ** state["step"]) else: - bias_correction = 1 + exp_avg_hat = exp_avg + exp_avg_sq_hat = exp_avg_sq - # Apply bias to lr to avoid broadcast. - step_size = group["lr"] * bias_correction + step_size = group["lr"] weight_norm = torch.norm(p.data).clamp(0, self.clamp_value) - adam_step = exp_avg / exp_avg_sq.sqrt().add(group["eps"]) + adam_step = exp_avg_hat / exp_avg_sq_hat.sqrt().add(group["eps"]) if group["weight_decay"] != 0: adam_step.add_(p.data, alpha=group["weight_decay"])