Handing the objective to autograd
Write only the objective in PyTorch — the model and the loss — and backward() produces the gradient you derived by hand. Store each factor as exp(θ): the model is identical, but there is no constraint to bump into and the gradient stays well scaled.
Adam at lr = 0.05 lands where L-BFGS-B did. Plain gradient descent diverges there and at 1e-4; it needs a step about five thousand times smaller (1e-5) and still finishes short, because the gradient scales differ wildly across the factor matrices — 7 days against 78 community areas.
Momentum (0.9) closes most of the gap; Adam’s per-parameter step size closes the rest. A differentiable objective can also be evaluated on a sample of entries — stochastic GCP, for tensors too large for memory.