PaperlayerView exact revision on arXiv

Source-linked AI summary

Local proxies for memory-efficient training

Wu S. Zihan, Brian A. Yuan, Guillaume Bellec

arXiv:2610.23145v1cs.LGcs.AI

TL;DR

Backpropagation’s activation storage makes training memory-intensive, while local learning has often traded accuracy for reduced memory and has seen limited success in transformers. ProxyProp distills a local proxy for each block to estimate global gradients. It is exact in deep linear networks, has bounded gradient error under stated residual-network conditions, and reports favorable accuracy–memory and LLM fine-tuning comparisons.

  • Problem

    Backpropagation requires storing activations until gradients propagate backward, and local learning has typically reduced memory at an accuracy cost while few algorithms have scaled successfully to transformers.

    Source: §1 · p. 1

  • Method

    ProxyProp partitions the transformer into blocks and distills a local proxy to predict the global output, allowing gradients to be propagated through the proxy instead of downstream blocks.

    Source: §1 · p. 1

  • Results

    ProxyProp offers a better accuracy/memory trade-off than other local algorithms on CIFAR-100 ViTs and higher performance than MeZO at comparable memory cost for LLM fine-tuning.

    Source: §1 · p. 2

  • Takeaways & Limitations

    ProxyProp’s activation-memory requirement is independent of network depth, unlike backpropagation’s depth-scaled activation storage.

    Source: §3 · p. 3

    The studied algorithm retains a pre-forward pass to compute the final activation and loss gradient, so it is not fully unlocked.

    Source: §3 · p. 3

Abstract

from arXiv · show

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.

1 INTRODUCTION

Backpropagation stores activations across network depth, while local learning can reduce that memory burden but has had limited success in transformers. ProxyProp uses block-local proxies to estimate global gradients, with theoretical guarantees and reported accuracy–memory benefits in vision transformers and language-model fine-tuning.

1 INTRODUCTION

  • Standard backpropagation stores activations until gradients return, so memory grows with depth; this can constrain LLM depth and sequence length under limited memory.The paper also identifies large memory use as a manufacturing cost and a source of datacenter training energy consumption.Source: §1 · p. 1
  • Local learning seeks to avoid backward locking, but few such algorithms have scaled to transformers and prior attempts have had limited success.The paper seeks a mathematically grounded, memory-efficient alternative without compromising training time.Source: §1 · p. 1
  • ProxyProp partitions transformers into gradient-decoupled blocks and distills a local proxy to predict the global output, then propagates each block’s gradients through that proxy.Local block updates can be independent, supporting efficient model parallelism.Source: §1 · pp. 1–2
  • A linear proxy implicitly estimates the expected network Jacobian, making ProxyProp exact for deep linear networks.The stated exactness applies to deep linear networks.Source: §1 · p. 2
  • For non-linear residual networks, a residual proxy gives exact local gradients at initialization and bounded estimation error when skip connections dominate.The bound is stated for residual contributions small relative to skip connections.Source: §1 · p. 2
  • Across various network partitionings in supervised classification with vision transformers, ProxyProp offers a better memory-accuracy trade-off than other local algorithms; for LLM fine-tuning, it achieves higher performance than MeZO at comparable memory cost.This LLM comparison holds even when MeZO uses 8x more compute time.Source: §1 · p. 2

2 REVIEW: STATE OF THE ART

Local learning has been pursued for different aims, including biologically plausible circuitry and decoupled training of network blocks. Prior methods use local supervised or self-supervised losses and auxiliary gradient predictors, but the latter can accumulate errors in practice.

2 REVIEW: STATE OF THE ART

  • Local-learning research spans biologically plausible circuitry and alternatives to symmetric feedback, while also complementing work on replacing backpropagation through time.The review connects this work to applications in physical systems, photonics, analog hardware, and digital ASIC technology.Source: §2 · pp. 2–3
  • Blockwise local learning stops gradients across blocks and trains each block with a supervised or self-supervised loss, including in pretrained-model fine-tuning.How stacks of such decoupled blocks build representation hierarchies remains unclear.Source: §2 · p. 3
  • Local gradient estimation seeks to approximate each block’s backpropagation gradient using local information.The paper contrasts this goal with prior local theories and auxiliary-network approaches.Source: §2 · p. 3
  • Auxiliary networks trained to predict gradients from the following block can become numerically challenging because their errors compound.The cited approach is Jaderberg et al. (2017).Source: §2 · p. 3

3 THEORY: PROXY DISTILLATION WITH ACCURATE GRADIENT PROPAGATION

ProxyProp replaces gradients through later network blocks with gradients through local proxies trained to predict the global output. Its linear proxy is exact for deep linear networks under stated conditions, while a residual proxy gives exact initial gradients and bounded error when skip connections dominate.

Definition of ProxyProp

  • ProxyProp trains a local proxy to predict the global network output by distillation, then propagates each block's loss gradients through that proxy instead of later, expensive layers.Distillation can align gradients even without explicitly minimizing gradient differences; under random perturbations of the block activation, minimizing the objective also minimizes squared Jacobian error.Source: §3 · pp. 3–4
  • ProxyProp releases each block's stored activation after its update, so activation memory is independent of network depth, unlike backpropagation.The algorithm's block updates are unlocked, although Algorithm 1 still requires a pre-forward pass to compute the final activation and loss gradient.Source: §3 · p. 3

Gradient consistency in linear proxies

  • For deep linear networks with Gaussian inputs of positive-definite covariance and full-row-rank weights, ProxyProp with an optimal linear proxy implements exactly the backpropagation algorithm.The optimal linear proxy's matrix is the expected Jacobian of the remaining layer stack.Source: §3 · pp. 4–5

Numerical verification in linear networks

  • For nonlinear networks, an optimal linear proxy can still yield imperfect gradient alignment because the expected Jacobian does not equal the sample-specific Jacobian.In the MNIST toy MLP, ProxyProp approaches the best linear projection achievable and has greater alignment than a random-projection DFA comparison.Source: §3 · p. 5

Gradient consistency in non-linear proxies

  • A transformer proxy should be a smaller nonlinear sequence model because token-wise linear proxies cannot track attention's temporal computation.Trainable linear projectors map the network's hidden dimension to the proxy's dimension.Source: §3 · p. 5

Gradient consistency in residual proxies

  • In nonlinear residual networks, a residual proxy gives exact gradients at initialization and bounded gradient-estimation error when residual Jacobians are small relative to skip connections.The bound assumes fixed depth and residual Jacobian operator norms at most ϵ for both the network and proxy.Source: §3 · p. 6

Numerical advantage of residual proxy

  • In a six-layer residual MLP, the residual proxy's gradients align perfectly at zeroed initialization and above 0.9 away from initialization; distillation improves alignment further.The residual proxy also aligns at the first iteration and converges faster in the Figure 2 illustration.Source: Figure 2 · p. 5, §3 · p. 6

4 ACCURACY/MEMORY TRADE-OFF IN TRANSFORMERS

On CIFAR-100 ViTs, ProxyProp offers a favorable accuracy–memory trade-off, with performance depending on proxy design and block partitioning. Proxy distillation performs comparably to a gradient-targeting loss without its significant improvement in accuracy.

4 ACCURACY/MEMORY TRADE-OFF IN TRANSFORMERS

  • ProxyProp (3 × 4) achieves performance close to backpropagation while reducing VRAM requirements by 62%.The experiment uses a 12-layer ViT with embedding dimension 128 on CIFAR-100 and dimension-64 transformer proxies.Source: §4 · p. 7
  • Identity proxies outperform fixed-random DFA and DRTP, while nonlinear transformer proxies outperform linear ones.The authors interpret these comparisons as supporting residual structure in local-gradient estimation.Source: §4 · p. 7
  • Replacing proxy distillation with a DNI loss brings no significant improvement, despite increasing memory through second-order derivatives.The auxiliary-MLP version of DNI is unstable with this ViT architecture.Source: §4 · p. 7
  • ProxyProp lies on the accuracy/memory Pareto front across block counts, despite storing some proxy activations to propagate gradients.The study compares 3, 6, or 12 blocks; changing block count navigates the performance–memory trade-off.Source: §4 · p. 7
  • For LLM fine-tuning, the study targets activation memory, alongside weight and optimizer memory that can be reduced with quantization and LoRA, respectively.The paper identifies consumer-grade GPUs as a practical setting for LLM fine-tuning.Source: §4 · p. 7

C Performance across models

Across three tasks and two model architectures, ProxyProp approaches backpropagation performance and exceeds reported MeZO performance. It uses less VRAM than backpropagation and trains faster than MeZO, while remaining roughly 1.5 times slower than backpropagation.

C Performance across models

  • MeZO reduces activation storage through repeated forward passes, but noisy gradient estimates can slow convergence and require more optimization steps.The paper describes this as the motivation for comparing ProxyProp with zeroth-order optimization.Source: D Wall clock time · p. 8
  • ProxyProp approaches backpropagation performance across three tasks and two architectures, and exceeds reported MeZO results.The comparisons use OPT-13B and LLaMA-3.2-3B on SST-2, WiC, and SQuAD; alternatives without residual parameterization and a blockwise local-loss baseline struggle to match it.Source: D Wall clock time · p. 8, Figure 4 · p. 8
  • ProxyProp is approximately 1.5 times slower than backpropagation on the same hardware, but roughly as fast overall.The comparison swaps PyTorch backpropagation for ProxyProp within the standard HuggingFace Trainer.Source: D Wall clock time · p. 8
  • ProxyProp uses 32 GB VRAM versus MeZO’s approximately 25.7 GB and backpropagation’s 64 GB; MeZO takes eight times the wall-clock time and still does not reach ProxyProp’s performance.The MeZO comparison uses 16 times as many steps as reported in Malladi et al. (2023).Source: Figure 4 · p. 8, D Wall clock time · p. 9
  • With four blocks, ProxyProp fits sequences over 50% longer than backpropagation; more blocks let it fit longer sequences and approach inference memory scaling.The comparison combines ProxyProp or backpropagation with LoRA, chunked loss, and gradient checkpointing.Source: D Wall clock time · p. 9

5 DISCUSSION

ProxyProp estimates global gradients locally and achieves near-backpropagation accuracy on ViTs and LLM fine-tuning. The discussion highlights potential memory benefits for deeper models and possible extensions to looped and diffusion architectures, while noting unresolved distributed implementation questions.

5 DISCUSSION

  • ProxyProp is presented as a local gradient-estimation method with theoretical guarantees and near-backpropagation accuracy in ViTs and LLM fine-tuning on a low-memory budget.The paper frames these outcomes as practical evidence following its theoretical analysis.Source: §5 · p. 9
  • ProxyProp’s memory use is constant with layer count, so its memory benefit grows with depth; scaling to much deeper architectures remains a potential application.The discussion contrasts this prospect with standard backpropagation’s difficulty training very deep architectures.Source: §5 · p. 9
  • For looped transformers, proxy propagation could estimate gradients aware of repeated loops without increasing memory; the paper also suggests an analogy between residual blocks and diffusion steps.The discussion proposes these as potential extensions, not evaluated results.Source: §5 · p. 9
  • Algorithm 1 retains weak coupling through its pre-forward pass; proposed routes to decoupling include accepting lagged loss gradients or using a fully unlocked variant.The authors had not studied distributed implementations at the time of writing.Source: §5 · pp. 9–10
  • The authors report that different authors trained models on different datasets with consistent speed and memory requirements, and that results were consistent across datasets.They state that datasets are publicly available and that implementation and a Hugging Face Trainer wrapper would be published upon acceptance.Source: REPRODUCIBILITY STATEMENT · p. 10

A SCALING OF MEMORY AND COMPUTE

This section concerns compute and memory costs for training an L-block network. For transformers, optimal activation checkpointing is not O(L), while common checkpointing stores layer outputs and recomputes attention and MLP activations, yielding O(L) activation storage.

A SCALING OF MEMORY AND COMPUTE

  • The table covers how compute and memory costs scale when training an L-block network.Its caption notes that optimal activation checkpointing for transformers is not O(L), while a common approach stores each layer's output and recomputes attention and MLP activations, resulting in O(L) activation storage.Source: Table A1 · p. 15
  • The proxy output pℓϕ(zℓ) is always the identity in the described algorithm.Source: A SCALING OF MEMORY AND COMPUTE · p. 15
  • Proxy parameters are pretrained using the distillation loss from Equation 2.The algorithm description identifies this as a step in fully unlocked ProxyProp.Source: A SCALING OF MEMORY AND COMPUTE · p. 15

B PROOF FOR PROPOSITION 3.4

The proof compares ProxyProp’s gradient estimate with backpropagation by expanding products of residual and identity contributions. Under the stated residual-map bounds, the estimation error is bounded to first order in ϵ.

B PROOF FOR PROPOSITION 3.4

  • The proof expands the product of Jacobians above block ℓ into paths choosing either a skip connection or a residual map at each layer.The all-identity path contributes I and is the contribution estimated by the identity proxy.Source: B PROOF FOR PROPOSITION 3.4 · p. 15
  • The gradient-estimation error is bounded by (L − ℓ + 1)ϵ + O(ϵ²) under the proof’s stated bounds.The constant C is defined using the block-parameter Jacobian and the loss gradient with respect to the final activation.Source: B PROOF FOR PROPOSITION 3.4 · pp. 15–16
  • The argument requires ∥Qℓ∥op ≤ ϵ, not a linear proxy, so the bound also applies to nonlinear proxies.Source: B PROOF FOR PROPOSITION 3.4 · pp. 15–16
  • The identity proxy, qℓϕ = 0 and Qℓ = 0, is a valid special case of the stated bound.Source: B PROOF FOR PROPOSITION 3.4 · p. 16

C EXPERIMENT DETAILS

The numerical verification uses a six-layer MNIST MLP to examine gradient alignment with proxy-based estimates. Figure 2 compares closed-form and trained proxies in plain or residual-network settings.

C.1 NUMERICAL VERIFICATION

  • The verification uses a six-layer MLP on MNIST with width 64 and batch size 4096.Each layer has a transformer-style branch with hidden width 256, tested with or without a skip connection.Source: §C.1 · p. 16
  • Figure 2A measures alignment between each layer’s local proxy gradient and its true gradient, using a linear proxy solved in closed form.The proxy targets zL for the plain network and zL − zℓ for the residual network.Source: §C.1 · p. 16
  • Figure 2B measures layer 0 alignment after training a linear proxy with Adam for 800 steps on a frozen random residual network.The flat line represents the closed-form solution.Source: §C.1 · p. 16

C.2.1 SPECIFICATIONS

The specifications describe CIFAR-100 vision-transformer experiments and memory profiling, followed by fine-tuning configurations for OPT and LLaMA models. The fine-tuning setup uses small proxies and shared classification heads across blocks.

C.2.1 SPECIFICATIONS

  • CIFAR-100 experiments use a 12-layer, width-128 vision transformer with four attention heads and patch size 4.The final mean-token readout uses a zero-initialized classifier, and the final LayerNorm is placed in the classifier so local methods read raw boundary activations.Source: §C.2.1 · p. 16
  • Training uses AdamW with batch size 128, learning rate 5 × 10−4, warm-up, cosine decay, and 100 epochs; peak memory is measured with batch size 512.ProxyProp uses a width-64 transformer proxy, and reported accuracies are from the last epoch.Source: §C.2.1 · p. 16

C.3 FINETUNING

  • Fine-tuning uses OPT-13B and LLaMA-3.2-3B on SST2, WIC, and SQuAD, with the models split into eight or seven blocks, respectively.Proxies are small transformer-based modules; all blocks share the classification head, and memory profiling uses synthetic data with sequence-chunked cross-entropy.Source: §C.3 · p. 17

D.1 VIT EXPERIMENTS

Extended ViT ablations compare proxy parameterizations and head placement, while SQuAD fine-tuning examines proxy training and loss-head placement.

D.1 VIT EXPERIMENTS

  • For unlocked updates, a 64-wide transformer proxy is trained on 20% of each epoch's batch with the backbone frozen, then frozen while the backbone trains on the full batch.Source: §D.1 · p. 17
  • A baseline that backpropagated classification loss through a block with an attached proxy suggests distilling the downstream network is more useful than merely adding a proxy to the local loss.Source: §D.1 · p. 17

D.2 FINETUNING

  • On SQuAD, identity proxies performed well, trained proxies improved performance, and prepending the loss head to the last block helped; OPT-13B runs used 8 blocks and proxy hidden dimension 80.The table caption reports 95% confidence intervals from five random seeds.Source: §D.2 · p. 17, Table A3 · p. 18

D.1 VIT EXPERIMENTS

  • ViT ablations found no significant classification difference between residual and non-residual proxies, but residual proxies had a clear fine-tuning advantage; prepending layers to the head improved performance at similar memory cost.Table A2 covers extended ViT ablations and baselines on CIFAR-100.Source: §D.1 · p. 17, Table A2 · p. 18
Loading 2610.23145v1…