m0 / v1 init

Allo @lucidrains , I’ve been fiddling with this optimizer, looking promising so far. I was looking for other interpretations out there for my doubts re no bias correction… I’m assuming it’s deemed unecessary due to the explicit m0 and v1 init, but wasn’t 100% sure it wasn’t just left out for clarity.

I noticed you left m0 as zero, and v1 as interpolation with zero init… did you experiment with that vs the notes in paper, Algorithm 1?

The core of my attempt below (note I flipped the betas to be comparable to adam/lamb/etc: .98, .92, .99)

    state = self.state[p]
    if len(state) == 0:
        state['step'] = 0
        state['grad'] = torch.zeros_like(grad)
        state['m'] = torch.clone(grad)  # init m0 = g0
        state['v'] = torch.zeros_like(grad)
        state['n'] = torch.zeros_like(grad)

    m, v, n = state['m'], state['v'], state['n']
    # NOTE first step is no-op as we need g0 & g1 for first grad delta (g1 - g0)
    if state['step'] > 0:
        m.lerp_(grad, 1. - beta1)
        grad_delta = grad - state['grad']
        if state['step'] > 1:
            v.lerp_(grad_delta, 1. - beta2)
        else:
            v.copy_(grad_delta)  # init v1 = g1 - g0
        n.lerp_((grad + beta2 * grad_delta).square(), 1. - beta3)

        # FIXME paper Algorithm 1 includes no bias correction
        # Does m0 and v1 init special cases obliviate the need or was left out of paper for clarity?
        denom = 1 + group['weight_decay'] * lr
        step_size = lr * (n + group['eps']).rsqrt()
        p.addcmul_(step_size, m.add(v, alpha=beta2), value=-1.).div_(denom)

    state['grad'].copy_(grad)
    state['step'] += 1

[XingyuXie]

Sorry for making something confused here. Adan indeed has the bias correction in the implementation, but we need to consist the algorithm presentation with the theoretical analysis. Hence, we did not explicitly emphasize it in Algorithm1. We’ll release the code in a few days (2-3 days since we have a code review procedure). The log and config files will release together. @rwightman

[XingyuXie]

@lucidrains Thanks for updating, the following are some minor modifications. When we implement Adan, we refer to some optimizer’s implementation in timm.

Line 55: state['prev_grad'] = grad Line 85-86:

correct_m = 1 / bias_correct1  # correction term for m'
correct_v = 1 / bias_correct2  # correction term for v

Line 91:

weighted_step_size = lr / ((n.sqrt()/sqrt_bias_correct3).add_(eps))

Tips:

• For fairness and ease of use, we do not enable the restart condition in practice.
• Adan can tolerate a large peak LR. For example, except for the experiments for the pre-training of MAE and LSTM, Adan’s LR is 5-10 times that of Adam/AdamW.
• Adan seems to be relatively sensitive to beta3. Adjusting beta1 and beta2 has a limited effect on the results, especially beta2.
• Interestingly, we found that weight_decay = 0.02 seems to be suitable for most experiments.