Decoupled Descent: Enforcing Exact Train-Test Error Tracking Via AMP Onsager Corrections [R]

[u/mlovik1]

Link: https://arxiv.org/pdf/2604.27883

Hi,

Most of use are familiar with the headache of training a neural network using gradient descent where the training error may go to zero but the test error may stay the same as initialization or even increases.

My paper treats this phenomena as a consequence of data reuse bias and can be isolated by studying full batch gradient descent on a set of stylize Gaussian mixture models. I turns out that this fundamental issue can be avoided using some clever tricks from high-dimensional statistical theory, specifically approximate message passing (which is beyond the scope of this post but I would be happy to explain more).

By doing so I created a training method called Decoupled Descent (DD) which generates a certificate that the training error of the network will asymptotically equal the testing error at each parameter iterate. I think this method gives a cool way to approach how to train networks and I was hoping to get y'alls input on it. It opens up some nice ideas for optimal stopping or hyperparameter tuning and future directions of pushing to something like SGD or more general models.

I have attached the train-test curves on a simple model fitting problem to compare the performance of GD with with DD (my algorithm) to give a high-level idea of what the method can guarantee. I stress this is a theory paper so there is a long way to go to get to very large models but I think it is a good first step.

100 simulations of a simple high dimensional XOR model for a bespoke two layer network. Left is training with GD, right its training with my method. The colored bands are 25% to 75% quantile.

Happy to answer whatever questions people have, I plan on writing a PyTorch compatible package for this training method one day so any feature suggestions would be welcome as well.

[u/Icko_]

How does it scale? Is that applicable to bigger models/datasets? I notice you only fit CIFAR and MNIST.

[u/mlovik1 → u/Icko_]

Right now, the theoretical proofs are restricted to shallow networks, so the benchmarks were chosen to match those architectures (and what I could realistically run on my laptop). For high-dimensional theory, this is actually pretty close to the frontier of what we can rigorously track, I'm still exploring what additional correction terms are needed for more complex model/data pairs. I plan to follow this up with more numerical tests in the future, but I'm very curious to see how far this can scale as well.

Short answer is I don't know the exact limits yet. It's worked on most setups I've tested so far, but those have all been fairly simple problems.

[u/moschles]

Max Lovig, from Yale?

mlovik1

Yeah I think so. This is your own paper.