Implementing WNGrad in Pytorch?
I'm trying to implement the WNGrad (technically WN-Adam, algorithm 4 in the paper) optimizier (WNGrad) in pytorch.  I've never implemented an optimizer in pytorch before so I don't know if I've done it correctly (I started from the adam implementation).  The optimizer does not make much progress and falls down like I would expect (bj values can only monotonically increase, which happens quickly so no progress is made) but I'm guessing I have a bug.  Standard optimizers (Adam, SGD) work fine on the same model I'm trying to optimize.

Does this implementation look correct?



from torch.optim import Optimizer

class WNAdam(Optimizer):
    """Implements WNAdam algorithm.

    It has been proposed in `WNGrad: Learn the Learning Rate in Gradient Descent`_.

    Arguments:
        params (iterable): iterable of parameters to optimize or dicts defining
            parameter groups
        lr (float, optional): learning rate (default: 0.1)
        beta1 (float, optional): exponential smoothing coefficient for gradient.
                                When beta=0 this implements WNGrad.

    .. _WNGrad\: Learn the Learning Rate in Gradient Descent:
        https://arxiv.org/abs/1803.02865
    """

    def __init__(self, params, lr=0.1, beta1=0.9):
        if not 0.0 &lt;= beta1 &lt; 1.0:
            raise ValueError("Invalid beta1 parameter: {}".format(beta1))
        defaults = dict(lr=lr, beta1=beta1)
        super().__init__(params, defaults)

    def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data
                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state['step'] = 0
                    # Exponential moving average of gradient values
                    state['exp_avg'] = torch.zeros_like(p.data)
                    # Learning rate adjustment
                    state['bj'] = 1.0

                exp_avg = state['exp_avg']
                beta1 = group['beta1']

                state['step'] += 1
                state['bj'] += (group['lr']**2)/(state['bj'])*grad.pow(2).sum()

                # update exponential moving average
                exp_avg.mul_(beta1).add_(1 - beta1, grad)

                bias_correction = 1 - beta1 ** state['step']
                p.data.sub_(group['lr'] / state['bj'] / bias_correction, exp_avg)
        return loss


The paper's author has an open sourced implementation on GitHub.
