sparse optimizers

[adamlerer]

Now that sparse tensors are mostly working (#1147), it would be awesome if all the optimizers worked with sparse tensors. This requires some cleverness to do amortized updates to the parameters. For example, for weight decay, you would do something like (pseudocode):

# apply weight decay
# dp is sparse
n_t++
for i in dp.indices():
  p[i] *= (1 - lr * wd)^(n_t - n[i])
  n[i] = n_t

Note this isn't exactly equivalent (to be equivalent you'd need to apply the weight decay before the forward pass, not after backwards), but it's a good approximation. You can do the same thing for momentum.

I'm guessing the same thing works for Adam/Adamax as well but I haven't worked through the equations. https://arxiv.org/pdf/1412.6980.pdf

@ezyang expressed interest in working on this.

[ezyang]

Here is an example test case for optimization when we have sparse gradients, done by "knocking out" the x or y coordinate on the computed gradient and putting the rest in a sparse gradient (I did it in a kind of janky way; ways to improve this code stylistically much appreciated!) When we get around to testing with weight decay I'll have to replace the assertEqual with a more approximate test, since the strategy we plan on using won't give us exactly the same results as the dense version.

    def _test_rosenbrock_sparse(self, constructor):
        params_t = torch.Tensor([1.5, 1.5])

        params   = Variable(torch.Tensor([1.5, 1.5]), requires_grad=True)
        params_c = Variable(torch.Tensor([1.5, 1.5]), requires_grad=True)
        optimizer   = constructor([params])
        optimizer_c = constructor([params_c])

        solution = torch.Tensor([1, 1])
        initial_dist = params.data.dist(solution)

        def eval(params, sparse_grad, w):
            optimizer.zero_grad()
            loss = rosenbrock(params)
            loss.backward()
            # params.grad.data now has the computed gradient. Turn
            # it into a sparse tensor
            params.grad.data.copy_(drosenbrock(params.data))
            # This is really goofy
            if w:
                i = torch.LongTensor([[0]])
                v = torch.DoubleTensor([params.grad.data[0]])
            else:
                i = torch.LongTensor([[1]])
                v = torch.DoubleTensor([params.grad.data[1]])
            x = sparse.DoubleTensor(i, v, torch.Size([2]))
            if sparse_grad:
                params.grad.data = x
            else:
                params.grad.data = x.to_dense()
            return loss

        for i in range(2000):
            w = torch.rand(1)[0] > 0.5
            optimizer.step(functools.partial(eval, params, True, w))
            optimizer_c.step(functools.partial(eval, params_c, False, w))
            self.assertEqual(params.data, params_c.data)

        self.assertLessEqual(params.data.dist(solution), initial_dist)

[ezyang]

Discussion with @adamlerer: if we amortize sparse updates in this way, we need some sort of "flushing" operation which a user can call to update the tensor to the "correct value" (since we may have accumulated updates on parameters that haven't been applied to the underlying tensor itself.) This is a change to the API for optimizers (a new flush function a user has to call at the end of optimizing) but since no one is actually optimizing on sparse models, this doesn't seem like a bad thing to ask them to do.

[ezyang]

Now that I feel that I have a much better grasp on what is going on with the optimizers, let me give an expanded "state of play":

• Adagrad is already converted to use sparse updates. https://github.com/pytorch/pytorch/pull/735/files#diff-533671a4db549248f40d40c6b5473916 This made the code a bit more convoluted, since the naive implementation converts the sparse tensors into their dense representations too early when it calls addcmul (it would be nice to make this nicer, though it is not obvious to me what strategy should be taken here.) However, adagrad is a "nice" algorithm, in the sense that when the gradient is sparse, the final update is sparse as well.

• Adam, however, the updates to the parameters are NOT sparse, even when the gradients are sparse, because it accumulates moments, which means a full update to the moments (to decay the old values) and the parameters (This was also observed by @martinraison). Indeed, TF has had to grapple with this issue RMSProp optimization support for sparse tensors tensorflow/tensorflow#464; at the time this bug was written, TF's Adam implementation did a dense update, causing performance problems for other users Slow Adam sparse updates in distributed TF tensorflow/tensorflow#6460. Eventually, to solve this problem, they merged a "lazy Adam optimizer" tensorflow/tensorflow@819a690d . The way this optimizer works is (1) only moments that show up in the gradient get updated, and (2) only those portions of the gradient get applied to the parameters.

  We could just implement this algorithm directly, but Adam's original suggestion about how to handle spare updates with weight decay got me thinking about how we might be able to play similar tricks in Adam's case, by applying the deferred momentum updates when we actually update the parameter in question (it's just a geometric series.) That's not a good idea: the point of momentum is to apply an extra force until the gradient tells us to stop. If we're training on sparse data and batch up momentum, you could easily shoot way past the optimal point before you got more information from gradient to go the other way. So we probably do want to go as TF goes.

• Reading about Nesterov, it seems to me that this trick should also apply when we are applying deferred updates from weight decay / momentum / etc. The big problem is that unlike with Nesterov, we don't know which deferred updates we are going to apply until we've computed our gradient once. If there was a way to quickly find out what the indices of the sparse update would be (without bothering to compute the entire gradient), we could first immediately apply any pending momentum/weight decay updates, and THEN compute the gradient on the new point (just like as with Nesterov.) I don't know enough about the autograd mechanism to know how easy or hard this would be to do.

• A simple matter: lazy weight decay should apply all optimizers. It will be annoying if we have to reimplement it each time, so a mixin is probably in order.

It will probably be interesting to see what TF is doing for all of the other optimizers we have.

[soumith]

We should probably introduce lazy adam. See https://twitter.com/haldaume3/status/901434195117051904 and https://discuss.pytorch.org/t/sparse-embedding-failing-with-adam-torch-cuda-sparse-floattensor-has-no-attribute-addcmul-/5589 as data-points

[jph00]

In the meantime, a documentation update would be helpful to document the sparse parameter to nn.Embedding and mention where it can't be used (weight decay and some optimizers).

[jph00]

I assume weight decay should work the same way as @ezyang points out lazy adam should be done. No reason to decay weights in batches that they're not, used, right?

[apaszke]

Well there might be a reason - if you do it lazily, the weight vector is never up to date, so the statistics you might compute (e.g. weight distribution) would be incorrect