Debiasing omission in Adam

A new paper: Revisiting Few-sample BERT Fine-tuning claims that bias in standard Adam makes the training process harder. (especially in low-data + fine-tuning)

https://www.youtube.com/watch?v=lWzo8CajF5s

So I wondered if debiasing had been implemented in the toolchains I was using.

1. Huggingface optimizer has an official option for debiasing Adam.
   

2. Pytorch also implemented debiasing in their Adam implementation.
   

https://github.com/pytorch/pytorch/blob/b162d95e461a5ea22f6840bf492a5dbb2ebbd151/torch/csrc/api/src/optim/adamw.cpp#L119

3. Nvidia FusedKernel doesn't implement 'debiasing'.
   

https://github.com/NVIDIA/apex/blob/082f999a6e18a3d02306e27482cc7486dab71a50/apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu#L67

I also found that Huggingface's implementation was just scaling step_size according to the bias, so the impact of the omission of debiasing might be negligible. (Someone might compensate the learning rate setting.) But Huggingface's simple scaling approach simply ignores the role of epsilon, I think PyTorch's implementation is more desirable.

Sketch code for the debias might be like below:

// these two coef's might be given as extra arguments.
const float bias_correction1 = 1 - pow(b1, step);
const float bias_correction2 = 1 - pow(b2, step);

for (int j = i; j < tsize; j+=totThreads) {
        T scaled_grad = g[j]/grad_scale;
        m[j] = b1*m[j] + (1-b1)*scaled_grad;
        v[j] = b2*v[j] + (1-b2)*scaled_grad*scaled_grad;
        
        float denom;
        if (mode == ADAM_MODE_0)
            denom = bias_correction1 * sqrtf(v[j]/bias_correction2 + eps);
        else // Mode 1
            denom = bias_correction1 * sqrtf(v[j]/bias_correction2) + eps;
        float update1 = (m[j]/denom) + (decay*p[j]);
        p[j] = p[j] - (step_size*update);
        if (p_copy != NULL) p_copy[j] = (GRAD_T) p[j];
}

Start with apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu at the linked line and compare its Adam update with the linked PyTorch implementation. Check how the kernel receives optimizer parameters and represents its two modes; done means the fused optimizer handles the requested bias correction consistently with the chosen semantics.