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.
Figures & tables
Figure 1: ProxyProp and its accuracy/memory trade-off. 1A : In ProxyProp, when updating block fℓ , only activations of that block need to be stored. The loss gradients are propagated through a proxy function instead of the other blocks. 1B : Memory and accuracy tradeoff for ProxyProp and other local methods; the 12-layer vision transformer (ViT) is either partitioned into 3, 6 or 12 blocks.
Figure 2: Numerical illustration of the theory. 2A : (Left) Alignment of the ProxyProp and backpropagation gradients in a random 6-layer MLP processing MNIST digits. (Right) Same as left, but the alignment is higher because of the skip connections. 2B : In the same residual setting, gradient alignment improves when the proxy objective is optimized (the network remains frozen). The residual proxy (solid) aligns already at the first iteration and converges faster.
Figure 3: Non-linear proxy with residual connection improves gradient alignment in trained transformers. Cosine similarity between the ProxyProp and the backpropagation gradients for the second block of a 12-layer ViT (partitioned into three blocks of four layers).
Method
top-1 accuracy
VRAM (GB)
3×4
6×2
12×1
3×4/6×2/12×1
ProxyProp
64.90±0.44
61.51±0.19
57.58±0.40
4.64 / 2.76 / 1.90
with DNI style loss
64.43±0.54
60.42±0.63
58.31±0.60
13.28 / 7.92 / 5.10
residual linear proxy
63.18±0.56
58.86±0.32
47.67±0.27
4.62 / 2.69 / 1.72
identity proxy
58.15±0.84
57.08±0.57
46.07±0.41
4.62 / 2.69 / 1.72
Backprop (one block)
63.93±0.60
60.00±0.54
50.33±0.32
4.44 / 2.47 / 1.49
Table 1: ProxyProp offers a better trade off between accuracy and memory cost . We report here the results on CIFAR-100 with a 12-layer ViT. We use a relatively large batch of 512 to highlight the memory gain of the activation vector. Our backpropagation baseline with the same architecture and training hyperparameters achieved 66.13% accuracy with 12.30 GB of peak VRAM usage. The accuracy number is in bold when it is highest or non-significantly lower (one-sided paired t -test over 5 seeds, Holm-corrected; p<0.05 ).
Figure 4: ProxyProp fine-tunes LLMs with better tradeoff between performance, memory, and compute time. 4A OPT-13B fine-tuned with rank-8 LoRA by Backprop, MeZO, and local algorithms; SST-2 and WiC (classification tasks) report accuracy, and SQuAD (question answering task) reports F1. MeZO values are taken from Malladi et al. (2023) ; Tan et al. (2025) . Error bars are the 95% confidence interval computed from 5 different random seeds. 4C Performance on OPT-13B with LoRA (solid) and LLaMA-3.2-3B (hatched). 4B , 4D Peak VRAM and compute time for one training step of input batch with size 8, sequence length 800 on A100 GPU. MeZO’s compute time is scaled up by its 16x number of training steps used in Malladi et al. (2023) .
Figure 5: ProxyProp approaches inference memory usage. We measure peak memory allocation when fine-tuning a LLaMA-8B model with varying-length sequence input. All baselines benefit from the standard memory-efficient fine-tuning implementation of the HuggingFace transformer library, including rank-8 LoRA, gradient checkpointing, and cross-entropy loss chunking. For ProxyProp we use a residual proxy transformer with 64 dimensions.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Method
Forward FLOPs
Backward FLOPs
Activation storage
Backprop (BP)
O(L)
O(L)
O(L)
BP + checkpointing
O(2L)
O(L)
O(L)∗
ProxyProp
O(2L)
O(L)
O(1)
Appendix
Table A1: Scaling of compute and memory cost with training an L -block network. ∗ For transformer, the optimal activation checkpointing is not O(L) because attention and MLP within each layer also requires plenty of memory storage. A common practice is to store the output of each transformer layer and recompute the activations of attention and MLP, which make the activation storage O(L) .
backprop
66.13±0.55
12.30
Method
top-1 accuracy
RAM (GiB)
3×4
6×2
12×1
3×4/6×2/12×1
ProxyProp
64.90±0.44
61.51±0.19
57.58±0.40
4.64 / 2.76 / 1.90
non-residual proxy
64.86±0.43
61.46±0.20
57.73±0.58
4.62 / 2.74 / 1.88
deep head
65.47±0.77
63.38±0.47
—
4.85 / 2.92 / —
fully-unlocked
64.78±0.42
59.79±0.20
42.02±1.28
4.60 / 2.67 / 1.71
Appendix
Table A2: Extended ablations and baselines of ViT on CIFAR-100.
Method
SQuAD Performance
Backprop
88.30±0.82
ProxyProp
84.75±0.91
ProxyProp with fixed identity proxy
82.15±1.82
ProxyProp + last block in the head
86.20±1.15
Appendix
Table A3: Performance of applying ProxyProp variants to finetuning OPT-13B (with LoRA-8) on SQuAD. All ProxyProp variants split the network into 8 blocks and use proxies with hidden dimension 80. The value after ± is the 95% confidence interval computed from 5 different random seeds.
The current deep learning training paradigm employs end-to-end backpropagation, regardless of the training stage, i.e. pre-training or fine-tuning. However, backpropagating through the entire model is neither biologically plausible nor memory efficient, since learning inside the brain is highly localized. Therefore, we propose LocalProp, a training procedure that locally updates the weights of a model. Our neuro-localized weight updates follow the "pre-training then fine-tuning" paradigm, where the pre-training is based on I-JEPA. After locally updating the weights, a pruning operation is performed, followed by a short final fine-tuning phase. Pruning helps by sending the learning signal from higher blocks to lower blocks. We perform experiments on several datasets, including large-scale benchmarks such as ImageNet, and empirically show that LocalProp reaches good performance at a fraction of GPU peak memory. By varying the number of jointly optimized blocks, we identify gradient-propagation span as a practical control over the accuracy-memory trade-off.
Diana-Nicoleta Grigore, Iuliana Georgescu, Radu Tudor Ionescu
Department of Computer Science University of Bucharest, Romania
LLM post-training typically propagates task gradients through the full depth of the model. Although this end-to-end structure is simple and general, it couples task adaptation to full-depth activation storage, long-range backward dependencies and direct task-gradient access to pretrained representations. We argue that this full-depth backward coupling can be unnecessarily expensive and intrusive, particularly when post-training supervision is much narrower than pre-training. To this end, we propose LoPT: Local-Learning Post-Training, a simple post-training strategy that makes gradient reach an explicit design choice. LoPT places a single gradient boundary at the transformer midpoint: the second-half block learns from the task objective, while the first-half block is updated by a lightweight feature-reconstruction objective to preserve useful representations and maintain interface compatibility. LoPT shortens the task-induced backward path while limiting direct interference from narrow task gradients on early-layer representations. Extensive experiments demonstrate that LoPT achieves competitive performance with lower memory cost, higher training efficiency and better retention of pretrained capabilities. Our code is available at: https://github.com/HumyuShi/LoPT
Hengyu Shi, Tianyang Han, Peizhe Wang +3
1Independent Researcher · 2D4 Lab · 3Southeast University
Large language models are trained with backpropagation, whose global gradient coordinates all layers but forces each to hold its activations and wait for the gradient to pass back through every deeper layer. Conventional local learning removes this update locking by training each module to predict the target through its own readout, but has not scaled to billion-parameter pretraining. We identify these private readouts as a key weakness, since they leave each module without information from deeper modules. We propose Shared-Output LOcal learning (SOLO), which replaces them with a shared, read-only copy of the final module's readout, the only one trained on the output of the whole network. Taken from the previous step, the copy transmits information from the final module without passing gradients between modules or reintroducing update locking. SOLO approaches backpropagation on Transformers of 340M to 2B parameters pretrained on 15B tokens, staying within one point in average zero-shot accuracy with a perplexity gap that narrows with scale. Readout ablations attribute SOLO's improvement over private readouts to sharing. Without update locking, each of p pipeline stages holds activations for O(1) micro-batches instead of O(p). The freed memory permits larger micro-batches, which reach up to 1.44x the best measured throughput of pipeline backpropagation on the same partition. To our knowledge, SOLO is the first local learning method to show such memory and throughput gains in billion-parameter language-model pretraining. Local learning thus becomes a practical alternative to backpropagation for large-scale pretraining.
Bojian Yin, Shurong Wang, Yuqi Pan +1
Institute of Automation, Chinese Academy of Sciences