Path-conditioned training: a principled way to rescale ReLU neural networks

Despite recent algorithmic advances, we still lack principled ways to leverage the well-documented rescaling symmetries in ReLU neural network parameters. While two properly rescaled weights implement the same function, the training dynamics can be dramatically different. To offer a fresh perspective on exploiting this phenomenon, we build on the recent path-lifting framework, which provides a compact factorization of ReLU networks. We introduce a geometrically motivated criterion to rescale neural network parameters which minimization leads to a conditioning strategy that aligns a kernel in the path-lifting space with a chosen reference. We derive an efficient algorithm to perform this alignment. In the context of random network initialization, we analyze how the architecture and the initialization scale jointly impact the output of the proposed method. Numerical experiments illustrate its potential to speed up training.

Training today's largest AI models is enormously expensive in both time and energy: every few weeks, AI labs around the world release a new language model, each one requiring months of computation and electricity consumption on the scale of hundreds of households over a year. Finding ways to accelerate this training would not only cut these costs, but also deepen our understanding of how machines actually learn. When you train an AI, you start with an "untrained" model that knows nothing and gradually shape it into a competent one. The starting point is typically chosen at random. We show that a more careful initial choice (derived from a theory that was missing until now) can roughly halve training time, while also explaining why the speed-up occurs. To make this idea easy to adopt, we have released a free, open-source tool. Anyone training an AI model can plug in our method and reduce the time and energy their experiments consume.