Local proxies for memory-efficient training

A paper posted as arXiv:2610.23145v1 proposes training a local proxy network for each transformer block to predict the global network output via distillation, letting gradients propagate independently through the proxy while approximating backpropagation gradients. The authors report the approach is exact for deep linear networks and that gradient error can be bounded in non-linear residual models with an appropriate residual proxy, and on CIFAR-100 vision transformers it achieves a better accuracy/memory trade-off than other local learning algorithms. The method also applies to fine-tuning large language models on a small memory budget, training faster and reaching higher performance than memory-efficient zeroth-order learning algorithms at comparable memory cost.

arXiv:2610.23145v1 Announce Type: cross Abstract: Local learning algorithms can reshape resource requirements of artificial intelligence. The memory bottleneck of standard end-to-end backpropagation lies in the need to store all network activations until the loss gradient has been propagated back through. Local learning avoids this, but typically at a cost in accuracy. We study how to achieve accurate and memory efficient local learning in transformers. For each transformer block, we train a local proxy network to predict the global network output via distillation. Gradients can be independently propagated through this local proxy while estimating well the backpropagation gradients. In theory, this approach is exact for deep linear networks; in non-linear residual models, the gradient error can be bounded when using the appropriate residual proxy. We verify this numerically and demonstrate that optimizing the distillation loss improves gradient estimation. Evaluating the performance of proxy propagation on CIFAR-100 in vision transformers, we find that it achieves a better accuracy/memory trade-off compared to other local algorithms. Finally, we show that the method is further relevant for fine-tuning large language models on a small memory budget. Our method trains faster and achieves higher performance than memory-efficient zeroth-order learning algorithms with comparable memory cost.