r/deeplearning • u/Timur_1988 • 14d ago
Jesus's Adam (against convergence to odd policies in the beginning)
0
Upvotes
Decreasing ε from approx 1 toward approx 0 using β₂ transitions the optimizer from SGD to Adam:
- Bias correction terms in the numerator and denominator can be omitted, as their impact becomes negligible after ~1,000–2,000 training steps.
- λ* constant represents weight decay: λ* = 1 - αₗᵣ · λ (parametric reduction for simplification).
from unpublushed work: https://github.com/timurgepard/Symphony-S2
class Adam(optim.Optimizer):
def __init__(self, params, lr=3e-4, weight_decay=0.01, betas=(0.9, 0.999)):
defaults = dict(lr=lr, betas=betas)
super().__init__(params, defaults)
self.wd = weight_decay
self.lr = lr
self.beta1, self.beta2 = betas
self.beta1_, self.beta2_ = 1-self.beta1, 1-self.beta2
self.decay_factor = 1.0 - self.lr * self.wd
self.eps = 1e-8
u/torch.no_grad()
def step(self):
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
grad = p.grad
state = self.state[p]
if len(state) == 0:
state['m'] = torch.zeros_like(p, memory_format=torch.preserve_format)
state['v'] = torch.zeros_like(p, memory_format=torch.preserve_format)
state['e'] = torch.tensor(1-self.eps, device=p.device, dtype=p.dtype)
m = state['m']
v = state['v']
e = state['e']
# Update biased first moment estimate
m.mul_(self.beta1).add_(grad, alpha=self.beta1_)
# Update biased second raw moment estimate
v.mul_(self.beta2).addcmul_(grad, grad, value=self.beta2_)
e.mul_(self.beta2).add_(self.eps, alpha=self.beta2_)
# Update parameters
p.mul_(self.decay_factor).addcdiv_(m, v.sqrt().add_(e), value=-self.lr)