PaperlayerView exact revision on arXiv

Source-linked AI summary

SOLO: Pretraining Billion-Parameter Language Models with Shared-Output Local Learning

Bojian Yin, Shurong Wang, Yuqi Pan, Guoqi Li

arXiv:2609.35440v1cs.LGcs.DC

TL;DR

Backpropagation coordinates language-model layers but makes them wait for gradients, while local learning removes that update locking yet has not scaled to billion-parameter pretraining. SOLO shares a read-only copy of the terminal readout across local heads. It approaches BP across tested scales, with saved activation memory supporting higher pipeline throughput.

  • Problem

    Local learning removes update locking but had not been shown to pretrain billion-parameter language models, where each module’s private readout also lacks information from deeper modules.

    Source: §1 · pp. 1–2

  • Method

    SOLO isolates gradients between modules while having every auxiliary head predict through a read-only copy of the terminal readout.

    Source: §1 · p. 2, §2.2 · p. 3

  • Results

    Across 340M–2B parameters, SOLO stays within one point of token-matched BP in average zero-shot accuracy and reaches up to 1.44× the best measured BP throughput.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Local learning becomes a practical option for billion-parameter pretraining when memory or interconnect bandwidth is limited.

    Source: §5 · p. 9

    Tests use two or four modules, leaving the quality-efficiency tradeoff for finer partitions open.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Backpropagation coordinates layers but creates update locking, while local learning removes that wait at the cost of global-gradient coordination. SOLO addresses the readout gap with a shared, read-only copy of the terminal readout, approaching BP while reducing activation memory and improving pipeline throughput.

Motivation

  • Conventional local learning had mainly approached BP on modest image-classification models, while language-model studies used small models, limited tokens, or post-pretraining objectives.The paper identifies the absence of billion-parameter language-model pretraining with local learning as the gap it addresses.Source: §1 · pp. 1–2
  • Private readouts leave heads without information from deeper modules; sharing supplies one consistent output space across modules, and ablations attribute SOLO’s gain over private readouts to sharing.The paper reports the readout-sharing effect in language modeling, but not image classification, where even a random readout costs little.Source: §1 · p. 2

Method

  • SOLO replaces each private auxiliary readout with a read-only copy of the terminal readout, which carries information from the final module without restoring update locking.Only the final objective updates the terminal readout; the previous-step copy is used by every auxiliary head.Source: §1 · p. 2, Figure 1 · p. 2

Systems

  • Without update locking, each of p pipeline stages holds activations for O(1) rather than O(p) micro-batches, freeing memory for larger micro-batches and higher throughput.The paper compares pipeline training with BP under the same partition into stages.Source: §1 · p. 2

Scale

  • Within one point of token-matched BP in average zero-shot accuracy, SOLO pretrained 340M–2B-parameter models; saved activation memory yielded up to 1.44× the best measured BP throughput.The runs used 15B tokens; SOLO is described as the first local-learning method to pretrain billion-parameter language models from scratch.Source: §1 · p. 2

2 Method

The method separates local gradient paths and residuals from the readout choice: SOLO keeps gradients isolated while sharing the terminal readout across auxiliary heads. This makes the readout factor match BP, although residual, path, temperature, and copy delay can still leave gradients misaligned.

2.1 Local learning

  • BP sends modules a gradient through the final readout, whereas conventional local learning cuts inter-module gradients and trains each module with its own auxiliary head and readout.The auxiliary heads are discarded after training.Source: §2.1 · p. 3

2.3 How sharing changes the local gradient

  • Sharing matches the readout factor of the local gradient to BP, while the path and residual remain local and do not require current-sample signals from deeper modules.The readout is a parameter rather than a signal computed on the current sample.Source: §2.1 · p. 3, §2 · p. 3

2.2 SOLO: local learning with a shared readout

  • SOLO gives every auxiliary head a read-only copy of the terminal readout, updated only by the final loss, so information from the final module reaches heads without gradients passing between modules.Each head trains its own mapping and temperature; the readout copy comes from an earlier step.Source: §2.2 · p. 3

2.3 How sharing changes the local gradient

  • SOLO’s gradient need not perfectly align with BP: temperature and copy delay can restore a readout difference, and residual and path differences remain.The equation-based comparison separates the gradient gap into readout, residual, and path terms.Source: §2.3 · pp. 3–4
  • Readout-source comparisons vary only the auxiliary-head readout; RAND, PRIV, SOLO, and pretrained-readout variants distinguish random, learned-private, shared-live, and diagnostic frozen choices.All variants learn the same temperature freedom; pretrained readouts are diagnostic probes, not training methods.Source: §2.3 · p. 4, Table 1 · p. 4

3 Experiments

The experiments compare SOLO with token-matched backpropagation (BP), isolate the effect of shared readouts, and measure pipeline-training trade-offs. Across model scales, readout sources, and pipeline settings, they assess language-model quality, why sharing helps, and how saved activation memory affects throughput.

3.1 SOLO pretrains models up to 2B near BP quality

  • With two modules, SOLO stays within 1.01, 0.68, and 0.64 WikiText perplexity of BP from 340M to 2B parameters, while average zero-shot accuracy stays within 0.5 points.With four modules, the perplexity differences are 3.15, 1.98, and 1.28, and average zero-shot accuracy stays within 0.9 points; the relative perplexity gap narrows with scale.Source: §3.1 · p. 5, Figure 2 · p. 5

3.2 SOLO’s gain over private readouts comes from sharing

  • Sharing lowers WikiText perplexity versus private readouts by 0.42 and 0.44 at 340M and 1.3B with two modules, and by 0.60 at 340M with four.A shared basis helps even with a random readout; rotating the trained readout separately for each module raises perplexity by up to 13.3 over SOLO.Source: §3.1 · p. 5, §3.2 · p. 6
  • In the one-epoch WikiText-103 sweep, SOLO leads PRIV by 2.1–4.7 perplexity from 40M to 512M, and the gain persists in 30-epoch Transformer-XL runs.The Transformer-XL gains are 0.15 and 0.42 perplexity with two and four modules, respectively.Source: §3.2 · p. 6
  • Sharing helps in language but not consistently in vision; a random readout costs up to 33% in language when vocabulary size is 26–85 times model width.In vision, a random readout costs little when classifier classes are fewer than feature dimensions.Source: §3.2 · p. 6

3.3 SOLO turns saved activation memory into pipeline throughput

  • At p=8, SOLO reduces peak activation memory from 17.4 to 2.5 GB; larger micro-batches reach 1.44× the best 1F1B throughput at micro-batch 16.The activation-memory reduction is 7.0-fold, close to p; SOLO reaches 1.43× throughput at equal memory.Source: §3.3 · p. 7
  • SOLO removes the pipeline backward bubble; with a 50-step refresh it reaches 1.18× 1F1B throughput at M=24, matches 1F1B at M=72, and beats VPP by 10% at M=24.Per-step refresh can cost up to 7%; a 10-step refresh recovers that throughput.Source: §3.3 · p. 7

3.4 Ablations

  • For refresh periods up to S=50, validation perplexity changes by at most 0.4%; S=100 and S=200 raise it by 1.9% and 3.0%.A 10–50-step period keeps quality while recovering the 7% throughput cost of per-step refresh.Source: §3.4 · p. 7
  • Deepening heads from H=0 to three blocks shrinks SOLO’s WikiText perplexity gap to BP from 8.6–13.5% to 2.7–5.2%; beyond H=2, added blocks narrow it by only about one point each.PRIV remains 5.4–7.8% behind SOLO across the sixteen cells, motivating the use of H=2.Source: §3.4 · p. 8

4 Related work

Prior local-learning work for language models remained limited in scale or training setting, while SOLO shares the terminal readout across modules. Other methods that reuse readouts or improve pipeline schedules retain end-to-end gradients or update locking.

Local learning

  • Language-model local-learning studies reached only smaller models, limited token budgets, repeated WikiText-103 passes, or post-pretraining objectives; SOLO instead shares the terminal readout across modules.Earlier examples include a 6M-parameter Transformer, 12-layer Llama-2-style models, and a 774M-parameter model trained over repeated WikiText-103 passes.Source: §4 · p. 8

Other ways to avoid update locking

  • Methods that avoid update locking pass downstream derivatives, estimates, targets, or final errors between modules, whereas SOLO passes only a parameter copy.Forward-only rules avoid derivatives but had not reached language-model pretraining in the cited work.Source: §4 · p. 8

Sharing a readout across depth

  • Early-exit and aligned models share or align readouts across depth but train end to end; SOLO updates its shared readout only with the final objective.Lenses decode intermediate layers after training, while fixed random classifiers suffice for image classification but not language in these experiments.Source: §4 · p. 8

Pipeline parallelism

  • Pipeline schedules reduce idle time but retain update locking, while SOLO eliminates the wait for gradients across depth.GPipe, 1F1B, and zero-bubble schedules still require stages to retain or recompute activations until gradients return.Source: §4 · p. 9

5 Discussion and Conclusion

SOLO shares one readout across modules while isolating their gradients, bringing local learning close to backpropagation and changing the memory and throughput trade-off. The evidence is bounded by the tested partitions, model architectures, and systems setting.

5 Discussion and Conclusion

  • Sharing a readout gives modules a common basis and information from the final module, bringing local gradients closer to backpropagation.The benefit is strongest when the readout is learned with local modules and the vocabulary is large, as in language.Source: §5 · p. 9
  • Experiments use two or four modules and plain Transformers; finer partitions, modern architectures, post-training, and convergence guarantees remain unaddressed.The systems study runs on one node, and on 24-layer pretraining models SOLO trails 1F1B at a fixed micro-batch.Source: §5 · p. 9
  • Future work includes cheaper heads for finer partitions, asynchronous readout refresh for pipelines spanning nodes, and extensions to modern architectures and post-training.The authors also suggest that an exponential moving average of the readout may bring the benefit of a mature copy into a single run.Source: §5 · p. 9
  • SOLO brings local learning close to backpropagation from 340M to 2B parameters, and each pipeline stage holds activations for O(1) micro-batches.The paper presents this as making local learning a practical pretraining option when memory or interconnect bandwidth is limited.Source: §5 · p. 9

A The SOLO step

SOLO updates each module after its own local passes rather than waiting for deeper modules. In pipeline mode, activations stream forward and the terminal readout is broadcast without passing gradients between modules.

A The SOLO step

  • Each module updates after its own forward and local backward passes, before the next module starts, so it does not wait for a deeper module.In a single process, the modules run in sequence.Source: A The SOLO step · p. 12
  • In one process, a stopped-gradient view of W supplies the previous-step readout copy at no cost; in a pipeline, the last GPU broadcasts W every S steps.The pipeline places each module on its own GPU and streams activations forward.Source: A The SOLO step · p. 12
  • No gradient passes between modules; Algorithm 1 marks this step's updated parameters blue and the read-only copy orange.The copy is not updated by local losses.Source: A The SOLO step · p. 12

B The local gradient when readout and path match

The analysis separates local-versus-backpropagation gradient differences into readout, residual, and path terms. Matching the readout and path gives a nonnegative gradient inner product, with equality when the residuals also match; a separate condition describes positive alignment more generally.

B The local gradient when readout and path match

  • Equation (4) decomposes the local-versus-backpropagation gradient gap into readout, residual, and path terms, then considers when readout and path terms vanish.The appendix gives a proposition for a head that predicts through τkWk.Source: B The local gradient when readout and… · p. 12
  • Matching τkWk = W and Jk = Mk gives a nonnegative inner product between local and backpropagation gradients; matching δk = δ as well makes the gradients equal.The equality follows when the readout, path, and residual terms in the gap decomposition vanish.Source: B The local gradient when readout and… · p. 12
  • The inner product can still be negative because a positive semidefinite form can be negative off its diagonal.Source: B The local gradient when readout and… · p. 12
  • Under cross-entropy, local-gradient alignment is positive when the head's distribution is sufficiently close to the final module's, but the path condition does not generally hold.The alignment condition is ∥δk − δ∥Gk < ∥δ∥Gk; a shallower head may not match the path through later modules.Source: B The local gradient when readout and… · p. 12

C Experimental details

The experimental details describe a relative-perplexity-gap comparison during language-model pretraining and vision sweeps across image-classification datasets and backbones.

C Experimental details

  • The 340M and 1.3B language-model curves track SOLO validation perplexity divided by BP perplexity at matching token counts, with parity marked at 1.The plotted ratios come from single-seed training-log curves, without interpolation or smoothing.Source: Figure 6 · p. 13
  • The ResNet-32 vision sweep covers CIFAR-10, CIFAR-100, and Tiny-ImageNet using fifteen modules and four-layer convolutional auxiliary heads.Each auxiliary head has a 64-dimensional readout.Source: C Experimental details · p. 13
  • Single-seed ResNet-50 scouting runs repeat the three vision datasets with readout width 2048, and cross-dataset comparisons are qualitative because schedules differ.Source: C Experimental details · p. 13

D Readout sweeps

The readout sweeps compare language and vision settings and test how auxiliary-head depth, readout sharing, and readout basis relate to performance. On the language sweep, the reported ablations show a consistent advantage for shared readouts across head depths.

D Readout sweeps

  • The auxiliary-head-depth sweep varies H for SOLO and PRIV at three language-model sizes, with H=8 included only at 256M as an over-deep control.H=0 is a linear readout; the Table 8 grid refreshes the shared copy every step, unlike Table 4’s 100-step refresh.Source: D Readout sweeps · p. 13, Table 8 · p. 14
  • Table 4 reports three-seed WikiText-103 validation perplexity across language-model sizes and identifies the best local-learning variant per column.Its two PRETRAINED rows are diagnostic probes using frozen terminal readouts, not training methods.Source: Table 4 · p. 14
  • The vision tables cover terminal test accuracy on three datasets for ResNet backbones and ImageNet-1k accuracy and memory change relative to each backbone’s end-to-end baseline.ResNet-50 vision columns are single-seed scouting runs; the ImageNet-1k comparison uses three-seed means.Source: Table 5 · p. 14, Table 6 · p. 14
  • The random-readout penalty is compared against output classes per unit width in width and vocabulary sweeps, which meet at V/d=64 with penalties of 29.80% and 29.72%.The difference is within run-to-run noise.Source: Table 7 · p. 15
  • 13.3 perplexity at H=0 and 7.4 at H=2 is the ROT cost over SOLO, while RSHARE beats RAND by 12.0 and 4.2, respectively.The PRIV–SOLO gap remains 5.5 and 5.3 at the two head depths; the reported effects exceed paired-seed standard deviation by more than ten times.Source: D Readout sweeps · p. 15
  • Under ROT at H=0, modules 2 and 3 have gradient alignment 0.67 and 0.61, versus 0.78 and 0.84 under SOLO.Two head blocks halve the last-module alignment gap from 0.23 to 0.11; the authors connect the mismatch to inherited basis and explain that skip connections and RMSNorm prevent exact reversal of Rk.Source: D Readout sweeps · p. 16

E Every depth of a SOLO model decodes

The section compares SOLO exits with the BP logit lens and examines what exits decode at each depth. SOLO exits remain close to the terminal model, while intermediate module outputs under SOLO are already readable by the terminal readout.

E Every depth of a SOLO model decodes

  • A SOLO exit uses the auxiliary head and shared readout it was trained with, whereas the BP logit lens applies the terminal norm and unembedding to an intermediate state never trained to be read.Source: E Every depth of a SOLO model decodes · p. 16
  • On the illustrated passage, the 6-layer exit reads most of the text; deeper exits change a handful of words, and the terminal differs from the end-to-end model at only a few positions.The changes most often involve words depending on an earlier mention.Source: E Every depth of a SOLO model decodes · p. 16, Figure 9 · p. 17
  • 71 to 75% of SOLO exits match the BP terminal’s top prediction on the battery, versus 16 to 51% for the BP logit lens.Every SOLO exit is within 2 perplexity points of its terminal, whereas the BP lens is not.Source: Table 10 · p. 17
  • Terminal-readout decoding gives SOLO intermediate module outputs about 0.6 to 0.75 top-1 agreement with BP from the first module onward, while PRIV, RAND, and PRETRAINED (BP) stay near 0.1 to 0.25 until the final module.All variants meet at the terminal module; the probe bypasses auxiliary heads and compares terminal-readout decoding with the matched-depth BP logit lens.Source: E Every depth of a SOLO model decodes · p. 17

F Refresh period of the shared readout

The refresh-period study measures how often the shared readout copy can be updated without materially changing validation perplexity. Periods up to 50 steps stay close to per-step reading, while longer periods incur larger reported costs.

F Refresh period of the shared readout

  • Up to S=50, the shared-readout copy stays within 0.3 perplexity of per-step reading; S=100 costs 1.9% and S=200 costs 3.0%.S=10 is marginally the best setting, and every exit moves together.Source: F Refresh period of the shared readout · p. 17, Table 11 · p. 17
  • After warmup, the readout moves about 1% of its norm over 10 steps, 3% over 50, 5% over 100, and 10% over 200.The study reports a 1.9% cost at 5% readout change and describes periods of 50 steps and below as harmless; the main pretraining runs use S=1.Source: F Refresh period of the shared readout · pp. 17–18
  • Figure 11 tracks relative terminal-readout change over refresh periods S∈{1, 10, 50, 100, 200}, with the shaded band marking learning-rate warmup and the final fall following cosine decay.Source: Figure 11 · p. 18

G Systems measurements

Across four A100s, SOLO trades throughput against lower memory than data-parallel backpropagation, with the throughput gap narrowing as model size grows. The section also compares off-NVLink behavior and pipeline backpropagation, and reports FLOPs and utilization measurements.

Against data-parallel backpropagation

  • Against FSDP, SOLO reaches 0.78 to 0.88 of throughput with K=2 from 340M up (0.86 at 2B), while using 11 to 25% less memory; with K=4 it reaches 0.64 to 0.74 of throughput and uses 35 to 48% less memory.At 2B, replicated DDP does not fit an 80 GB device.Source: Figure 12 · p. 19, G Systems measurements · p. 19

Setup

  • With each configuration at its own memory-optimal micro-batch, the 2B K=2 pipeline runs at 53.4k tokens per second in 42.9 GB, versus DDP at 46.6k in 55.2 GB.The K=4 pipeline runs at 44.9k tokens per second in 26.7 GB.Source: G Systems measurements · p. 19

Two baselines at 1.3B

  • At 1.3B and micro-batch 8, the K=2 pipeline saves 53% of memory versus DDP at 0.88 of its throughput, while K=4 saves 70% at 0.76.Against FSDP at 1.3B, K=2 saves 16% of memory at 0.78 of throughput, and K=4 saves 46% at 0.68.Source: G Systems measurements · p. 19

Interconnect

  • Off NVLink, DDP loses 72% of throughput while the pipeline loses 8%, reversing the ratio to 2.2 in this single-node emulation.The passage attributes the difference to DDP's synchronized, full-gradient all-reduce versus the pipeline's overlapped activation transfers.Source: G Systems measurements · p. 19, Table 12 · p. 20

FLOPs and utilization

  • Auxiliary blocks, not readouts, account for most of the training-FLOPs difference between BP and SOLO, and the speed denominator is plain BP.The comparison does not use an implementation that hides communication behind compute.Source: Table 13 · p. 20, G Systems measurements · p. 20
  • For 340M, ρ is 1.089, 1.266, and 1.621 at K=2, 4, and 8 with H=1, and 1.127, 1.382, and 1.891 with H=2.The passage notes that ρ is sensitive to head depth and module count.Source: G Systems measurements · p. 20

What is not measured

  • At 2B, pipeline backpropagation trades the memory cliff for 1F1B bubbles; its same-split comparison with SOLO uses two, six, and eight stages on a different model shape.The FSDP comparison off NVLink was not run; its predicted larger inversion is based on parameter-sized all-gather and reduce-scatter traffic.Source: G Systems measurements · pp. 19–20

H Comparison with pipeline-parallel backpropagation

Under the same pipeline split, SOLO avoids update waiting and reduces activation memory, but its throughput advantage depends on micro-batch size and model shape. At a fixed global batch, larger micro-batches let SOLO reach higher throughput than 1F1B, while shallow models split across many stages can incur greater per-module costs.

Throughput comparison

  • At p=8 and M=24, SOLO with S=10 is 10% faster than the fastest pipeline schedule in each cell, but an ideal interleaved schedule would match SOLO at M=24.The reported VPP implementation reaches 0.95 of its bubble-model throughput at v=2 and 0.88 at v=4 for M=24.Source: H Comparison with pipeline-parallel… · p. 22

Large M and a second model shape

  • For the 24-block pretraining model, SOLO reaches 0.82 of 1F1B throughput at M=24 and 0.72 at M=72; a one-block head at p=4 raises the ratio to 0.98.The per-module cost is higher for this shallow model than for the 96-block model, and the comparison’s sign depends on blocks per stage relative to head size.Source: H Comparison with pipeline-parallel… · p. 22
  • With the 96-block model at p=8, SOLO reaches 0.966 of 1F1B throughput at M=128 and 0.943 at M=256 when the readout copy refreshes every 50 steps.These ratios follow the fitted cost curve and approach 1/ρ = 0.92 from above.Source: H Comparison with pipeline-parallel… · p. 22

Memory

Micro-batch sweep

Limitations

  • The measurements are limited to one node with NVLink, synthetic tokens, and configuration (A); zero-bubble schedules were not measured, and the appendix reports systems costs only.The bubble model predicts ZB-H1 would be about level with SOLO at M=24 and ahead at M=72.Source: H Comparison with pipeline-parallel… · pp. 25–26
Loading 2609.35440v1…