WAME Algorithm Implementation Not Working

martinoywa: <p>Hi,<br>
I’ve tried implementing the Weight–wise adaptive learning rates with moving average estimator, WAME, based on this <a href="https://www.elen.ucl.ac.be/Proceedings/esann/esannpdf/es2017-50.pdf" rel="noopener nofollow ugc">paper</a>, but my loss doesn’t seem to be going down fast enough.<br>
Here’s my implementation:</p>
<pre><code class="lang-auto">import torch
from torch.optim import Optimizer


class WAME(Optimizer):
    """Implements the Weight–wise Adaptive learning rates with Moving average
    Estimator.

    Arguments:
        params (iterable): iterable of paramenters to optimize
        alpha (float, optional): smoothing constant (default: 0.9)
        etas (Tuple[float, float], optional): pair of (etaminus, etaplus), that
            are multiplicative increase and decrease factors
            (default: (0.1, 1.2))
        step_sizes (Tuple[float, float], optional): a pair of minimal and
            maximal allowed step sizes (default: (0.01, 100))
    """

    def __init__(self, params, alpha=0.9, etas=(0.1, 1.2), step_sizes=(0.01, 100)):
        if not 0.0 &lt;= alpha:
            raise ValueError(f"Invalid theta: {alpha}")
        if not 0.0 &lt; etas[0] &lt; 1.0 &lt; etas[1]:
            raise ValueError(f"Invalid eta values: {etas[0]}, {etas[1]}")

        defaults = dict(alpha=alpha, etas=etas, step_sizes=step_sizes)
        super(WAME, self).__init__(params, defaults)

    def __setstate__(self, state):
         super(WAME, self).__setstate__(state)


    @torch.no_grad()
    def step(self, closure=None):
        """Perform 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:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError('WAME does not support sparse gradients')
                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state["step"] = 0
                    state["prev"] = torch.zeros_like(p.data)
                    state["theta"] = torch.zeros_like(p.data)
                    state["z"] = torch.zeros_like(p.data)
                    state["step_size"] = grad.new().resize_as_(grad).fill_(1)

                etaminus, etaplus = group["etas"]
                step_size_min, step_size_max = group["step_sizes"]
                alpha = group["alpha"]
                step_size = state["step_size"]
                theta = state["theta"]
                z = state["z"]

                #print(state)

                state["step"] += 1

                mul_dx = grad.mul(state["prev"])

                if mul_dx &gt; 0:
                    step_size = min(step_size*etaplus, step_size_max)
                elif mul_dx &lt; 0:
                    step_size = max(step_size*etaminus, step_size_min)

                z = alpha * z + (1 - alpha) * step_size
                theta = alpha * theta + (1 - alpha) * (grad ** 2)

                new_grad = (z * grad) * (1 / theta)

                # update paramenters
                p.sub_(new_grad)

                state["prev"].copy_(grad)

        return loss

</code></pre>
<p>The algorithm itself is as below:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/0/f0a85f6c81779696d864fd6d7bcdbdc5e89c3ba9.png" data-download-href="https://discuss.pytorch.org/uploads/default/f0a85f6c81779696d864fd6d7bcdbdc5e89c3ba9" title="WAME"><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/0/f0a85f6c81779696d864fd6d7bcdbdc5e89c3ba9.png" alt="WAME" data-base62-sha1="ykXkrYNH7v0TtU6dIhMUdmBw27v" width="690" height="383" data-dominant-color="F6F6F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">WAME</span><span class="informations">936×520 25.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Any help as to what I’ve done wrong?</p>

martinoywa: <p>This is its description.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/5/258537f20cdc073c4772f295e125e2a119c72d6f.png" data-download-href="https://discuss.pytorch.org/uploads/default/258537f20cdc073c4772f295e125e2a119c72d6f" title="DESC"><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/5/258537f20cdc073c4772f295e125e2a119c72d6f.png" alt="DESC" data-base62-sha1="5lV5jPZPDr18wj4W4qXkArUDmSr" width="690" height="432" data-dominant-color="E9E9E9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">DESC</span><span class="informations">924×579 40.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>