PaperlayerView exact revision on arXiv

Source-linked AI summary

Greedy Local Learning for Language Model Pretraining: Gaps and Objective Design

Jihwan Moon, Sheir A. Zaheer, Jinmyoung Lee, Gunhee Kim, Chan Y. Park

arXiv:2610.04867v1cs.LGcs.AI

TL;DR

The loss gap between local learning and end-to-end training in larger-scale autoregressive LM pretraining, and auxiliary designs that might reduce it, had not been measured. This study compares block counts, model scales, and auxiliary designs at Chinchilla token budgets; the gap varies with block count and scale, and selected interventions improve the outcomes.

  • Problem

    How the likelihood gap of local learning behaves in autoregressive LM pretraining at 125M–400M parameters, and which auxiliary designs reduce it, had not been measured.

    Source: §1 · p. 2

  • Method

    The study trains 125M- and 400M-parameter Transformers with K ∈ {1, 2, 4} at Chinchilla token budgets, varying auxiliary network architecture and objective.

    Source: §1 · p. 2

  • Results

    The gap more than doubles from two to four blocks but shrinks with scale at K=4; first-block MTP is best at K=4, and per-block execution reduces activation memory by up to 2.2×.

    Source: §5 · p. 6

  • Takeaways & Limitations

    The measurements suggest local objectives should apply future-predictive pressure selectively across boundaries while resisting shortcuts that bypass predictive content.

    Source: §5 · p. 6

    The study is capped at 400M parameters, uses one seed per configuration, and compares token-matched rather than compute-matched runs; it offers measurements, not a concrete method.

    Source: §5 · p. 6

Abstract

from arXiv · show

Greedy block-wise local learning splits a network into gradient-isolated blocks trained by local auxiliary losses, deleting the backward pass between blocks: inter-stage communication becomes forward-only and every block can step its optimizer independently, properties directly relevant to decentralized model-parallel training. Local learning is competitive with end-to-end backpropagation on image classification, and on small Transformers it is known to trade a worse best loss for parallel speedup. How this loss gap behaves in autoregressive language model (LM) pretraining at larger scale, and which auxiliary designs reduce it, has not been measured. We present a token-budget-matched empirical study at 125M and 400M parameters with $K \in \{1,2,4\}$ blocks at Chinchilla-optimal budgets, factorizing the auxiliary design into network architecture and training objective. We observe: (i) the gap to end-to-end training more than doubles from $K=2$ to $K=4$, but at $K=4$ shrinks from 125M to 400M; (ii) replacing an MLP auxiliary with a Transformer-based one is a strong network-side intervention, recovering 22-41% of the gap; (iii) a multi-token-prediction (MTP) auxiliary objective helps at the first block boundary, whereas adding it at deeper boundaries hurts, and restricting it to the first block yields the best $K=4$ configuration ($+0.062$ vs. $+0.075$ nats at 400M); and (iv) deployment-style per-block execution reduces activation memory by up to $2.2\times$. We frame these results as an empirically grounded method direction rather than a finalized method: local objectives should apply future-predictive pressure selectively across boundaries while resisting shortcuts that bypass predictive content.

1 Introduction

Greedy block-wise local learning offers forward-only communication and independent block updates for decentralized model-parallel training, but its likelihood gap in larger-scale LM pretraining and ways to reduce it had not been measured. This study evaluates that gap across model scale and block count while factorizing auxiliary network and objective design.

Motivation and approach

  • Greedy block-wise local learning isolates K blocks with local auxiliary losses, preserving forward activations while removing cross-boundary activation gradients and backward update dependencies.In stage-partitioned deployment, inter-stage traffic becomes forward-only and each worker’s backward graph is limited to one block and its auxiliary head.Source: §1 · p. 2

Study scope

  • The study measures token-budget-matched LM likelihood gaps at 125M–400M parameters and Chinchilla-optimal budgets, extending small-Transformer results to larger-scale pretraining.It varies block count and auxiliary design.Source: §1 · p. 2

Per-block evidence

  • Per-block losses show later blocks barely improve inherited representations under local NTP supervision, consistent with representation myopia.The collapsed losses are evidence consistent with locally useful features failing to improve downstream, not direct proof that information was discarded.Source: §1 · p. 2

Auxiliary design

  • A Transformer-based auxiliary is a strong network-side intervention, while MTP helps at the first boundary but hurts when added at deeper boundaries.The comparison motivates boundary-adaptive, shortcut-resistant local objectives.Source: §1 · p. 2

Study contributions

  • Deployment-style per-block memory measurements and analytical accounting of communication and auxiliary training compute assess local learning as decentralized model parallelism.The factorized implementations used for the comparisons are released.Source: §1 · p. 2

2 Related work

Prior local-learning work was developed mainly for image classification, while the closest language-modeling study used small Transformers and a different emphasis. This study places local learning on the decentralized model-parallel axis.

Local learning

  • Image-classification work established competitive locally trained CNNs with DGL [Belilovsky et al., 2019, 2020], searched decoupling points and auxiliary architectures with SEDONA, and added reconstruction pressure with InfoPro [Wang et al., 2021].MAN and AugLocal incorporate later computation into local paths, introducing synchronization or replicated downstream computation in decentralized pipelines [Su et al., 2024, Ma et al., 2024].Source: §2 · p. 2

Language models

  • Laskin et al. [2021] studied local parallelism in small Transformer language models and compute–walltime trade-offs, whereas this study measures likelihood gaps across scale and auxiliary design at Chinchilla token budgets.LoPT locally post-trains pretrained models, while DiffusionBlocks replaces likelihood training with per-block denoising [Shi et al., 2026; Shing et al., 2026].Source: §2 · pp. 2–3

Decentralized training

  • Decentralized training includes data-parallel methods that reduce synchronization between replicas and model-parallel methods that distribute one model across weakly connected workers; this study concerns the latter axis.Local losses remove backward boundary traffic and update dependencies and can compose with data-parallel methods.Source: §2 · p. 3

3 Study design

The study partitions decoder-only Transformers into gradient-isolated blocks and varies each block’s auxiliary network and objective. It compares model sizes and block counts using fixed Chinchilla token budgets and held-out validation NLL.

3.1 Greedy block-wise local learning

  • Stop-gradient boundaries preserve the original forward activations while preventing gradients from crossing blocks; each non-final block trains from its own local auxiliary loss, and K=1 is standard end-to-end training.In pipelined deployment, workers receive forward activations and update from their local losses without backward update dependencies.Source: §3.1 · p. 3

3.2 Auxiliary networks: MLP-aux and Transformer-aux

  • With the local objective fixed to NTP, MLP-aux uses a pre-normalized SwiGLU FFN, while Transformer-aux adds a Transformer block with local self-attention and depth.Both heads are untied across blocks and discarded at inference.Source: §3.2 · p. 3

3.3 Auxiliary objective: multi-token prediction

  • MTP auxiliaries add offset-specific residual FFN branches predicting the next and second-next tokens; first-block MTP applies two-offset prediction only at the first auxiliary block, while all-block MTP applies a depth-weighted head at every auxiliary block.The final block remains NTP-only, and shifted targets add no inter-stage communication.Source: §3.3 · p. 3

3.4 Experimental setting

  • Experiments use 125M- and 400M-class Llama-style models on packed FineWeb-Edu sequences, with K ∈ {1, 2, 4} and strict Chinchilla-optimal token budgets.Held-out validation NLL per token is computed through the final block’s real NTP head; Table 1 reports absolute NLL for K=1 and deltas against it for K=2 and K=4.Source: §3.4 · p. 4, Table 1 · p. 4

4 Observations

The isolation gap grows with block count, while a Transformer auxiliary and first-boundary MTP reduce it. Per-block losses suggest representation myopia; local execution lowers activation memory but adds training compute.

4.1 Observation 1: the isolation gap grows with block count

  • For every auxiliary design, the gap more than doubles from K=2 to K=4; at K=4 it shrinks from 125M to 400M across all three designs.For example, first-block MTP falls from +0.074 to +0.062, while the lowest-gap operating point is +0.027 nats for first-block MTP at K=2, 400M.Source: §4.1 · p. 4

4.2 Observation 2: local losses exhibit representation myopia

  • At 125M and K=4, later blocks barely improve on inherited representations under local NTP supervision, a pattern consistent with representation myopia.The same qualitative pattern holds at 400M; the curves collapse even with a larger auxiliary head.Source: §4.2 · p. 4, Figure 1 · p. 5

4.3 Observation 3: the Transformer-based auxiliary helps most

  • Replacing MLP-aux with Transformer-aux reduces the gap by 39% and 41% at 125M (K=2, 4), and by 22% and 33% at 400M.The curves move downward without separating; because the heads are not parameter- or compute-matched, this is an architecture-level comparison, not a controlled estimate of attention alone.Source: §4.3 · p. 5

4.4 Observation 4: MTP depends on boundary position

  • MTP helps at the first boundary but hurts when applied at deeper boundaries; at K=4, first-block MTP improves over all-block MTP and Transformer-aux.For all-block MTP, next-token CE improves from 3.087 to 2.858 as the second-token weight decreases across blocks; the NTP+CPC hybrid instead increased the gap to +0.154 nats at 125M, K=2.Source: §4.4 · p. 5, Table 2 · p. 5, Figure 2 · p. 6

4.5 Observation 5: local execution reduces memory at a training-compute cost

  • At K=4, local execution cuts 400M peak memory from 29.2 to 14.5–15.8 GiB and activation memory from 27.7 to 12.4–13.4 GiB, up to 2.2×.The reduction is sublinear in K because each worker still runs a vocabulary-sized auxiliary head. Auxiliary heads add training compute: at K=4, Transformer-aux and first-block MTP total 78% and 97% of end-to-end training FLOPs at 125M, and 39% and 48% at 400M.Source: §4.5 · pp. 5–6, Table 4 · p. 6

5 Conclusion

Local-learning LM pretraining shows a growing loss gap with more blocks, while auxiliary design and scale affect that gap; per-block execution also reduces activation memory. The study remains bounded by its scale, seed count, and token-matched comparisons.

5 Conclusion

  • At K=4, the gap shrinks from 125M to 400M, while increasing blocks from K=2 to K=4 more than doubles it.The conclusion also reports that collapsed per-block losses are consistent with representation myopia.Source: §5 · p. 6
  • A Transformer-based auxiliary is an effective network-side intervention, and MTP helps at the first boundary but hurts at deeper ones.First-block MTP is the best K=4 configuration at both scales.Source: §5 · p. 6
  • Per-block execution reduces activation memory by up to 2.2×, at an added training-compute cost.Source: §5 · p. 6
  • The study is capped at 400M parameters, uses one seed per configuration, and compares token-matched rather than compute-matched runs.It offers measurements rather than a finalized method.Source: §5 · p. 6
  • The authors suggest local objectives apply future-predictive pressure selectively across boundaries while resisting shortcuts that bypass predictive content.Source: §5 · p. 6

A.1 Greedy block-wise local learning

Prior work established greedy local learning mainly in image classification and small language models, while related approaches to fine-tuning and decentralized training address different parts of the problem. This study focuses on local learning across model-parallel stages.

A.1 Greedy block-wise local learning

  • Earlier work studied locally supervised CNNs, information-preserving auxiliaries, and small-Transformer local updates, which traded best loss against compute–walltime speedups.For the small Transformer, Laskin et al. [2021] report best validation losses of 3.62 for backpropagation and 3.66 for overlapping updates; this study varies auxiliary design at larger scale instead of measuring compute–walltime frontiers.Source: §A.1 · pp. 8–9

A.2 Layer-wise fine-tuning of LLMs

  • Layer-wise fine-tuning methods select parameters under the end-to-end loss but still traverse every layer and adapt pretrained models rather than pretraining from scratch.LIFT and LISA [Pan et al., 2024] reduce fine-tuning memory, not inter-stage communication or update locking.Source: §A.2 · p. 9

A.3 Decentralized LLM pretraining

  • Decentralized pretraining includes data-parallel local-update methods and model-parallel methods that distribute a model across weakly connected workers or compress boundary traffic.Local learning complements model-parallel approaches by deleting backward traffic and update dependencies; the two axes can compose.Source: §A.3 · p. 9

B Training and implementation details

The implementation uses step-matched AdamW runs, while diagnostics and execution analyses examine local losses, MTP weighting, and the remaining pipeline dependencies.

B Training and implementation details

  • Runs use AdamW with a 3 × 10−3 learning rate, 0.033 weight decay, 1.0 gradient clipping, and bf16 mixed precision, with evaluations at fixed final steps.The final checkpoints are step 5,000 at 125M and step 14,000 at 400M, making comparisons step-matched.Source: B Training and implementation details · p. 9
  • The all-block MTP linear schedule places more second-token weight on earlier blocks; uniform weighting assigns 1/2 per offset at each auxiliary block.At K=2, first-block MTP and both all-block weighting schemes coincide at 1/2 per offset.Source: B Training and implementation details · p. 10
  • At K=4, separately executed stages peak at 10.14/10.21/10.32/9.36 GiB for Transformer-aux and 10.34/10.25/10.36/9.41 GiB for first-block MTP at 125M.At 400M, the corresponding peaks are 15.24/15.43/15.67/14.47 GiB and 15.55/15.51/15.76/14.56 GiB; sequential MTP execution keeps only one vocabulary-sized logits tensor live at a time.Source: B Training and implementation details · p. 10

C Per-block diagnostics at 400M

D Backward and update locking

  • Pipeline backpropagation requires each stage to await gradients from downstream modules, whereas local losses remove backward and update locking but retain forward dependency.Stages can stream successive microbatches through the remaining forward dependency, though pipeline fill, drain, and load imbalance remain.Source: D Backward and update locking · p. 10

E Analytical inter-stage communication

Local learning removes backward inter-stage traffic, halving the analytical per-boundary communication relative to pipeline backpropagation; MTP adds no inter-stage traffic.

E Analytical inter-stage communication

F Training-compute overhead of the auxiliary choices

Auxiliary heads add training compute, chiefly through each head’s vocabulary projection, while the relative overhead falls at larger model scale; these token-matched quality comparisons do not establish wall-clock savings.

F Training-compute overhead of the auxiliary choices

  • 16.0% of end-to-end FLOPs per head at 125M and 7.1% at 400M come from the vocabulary projection required for each local next-token loss.These percentages use analytical training-FLOP accounting at sequence length 4,096.Source: F Training-compute overhead of the… · p. 11, Table 6 · p. 11
  • A second vocabulary-projection pass adds 16.0 points at 125M and 7.1 points at 400M per MTP head, the main added cost of MTP.First-block MTP pays this cost at one head only.Source: F Training-compute overhead of the… · p. 11, Table 6 · p. 11
  • At 125M, K=4, a Transformer-aux head costs 0.24 GFLOP/token versus 0.19 for its trained block; at 400M, the costs are 0.35 versus 0.64.With fixed vocabulary size, overhead roughly halves from 125M to 400M and grows linearly with the number of auxiliary blocks, K−1.Source: F Training-compute overhead of the… · p. 11
  • Quality gaps are token-matched, not compute-matched, and the study does not measure throughput, so wall-clock gains remain deployment-dependent.Local updates remove pipeline bubbles and backward synchronization, but auxiliary compute is not charged in the quality comparisons.Source: F Training-compute overhead of the… · p. 11

G MTP placement and weighting ablation

At K=4, where boundary placement can vary, moving MTP pressure to the first boundary improves validation NLL more than changing its weighting schedule; K=2 cannot distinguish these placements.

G MTP placement and weighting ablation

  • 0.008 nats at 125M and 0.024 nats at 400M are the validation-NLL improvements from moving MTP pressure from all boundaries to the first boundary.The ablation compares K=4 configurations against Transformer-aux as the NTP-only reference.Source: G MTP placement and weighting ablation · p. 12, Table 7 · p. 12
  • ∼0.003 nats separates the linear and uniform all-block MTP schedules, at the level of run-to-run noise; the all-block uniform 400M configuration was not run.At K=2, first-block MTP, all-block linear, and all-block uniform coincide because there is one auxiliary block.Source: G MTP placement and weighting ablation · p. 12, Table 7 · p. 12
  • +0.017 versus +0.020 nats at 125M for the two K=2 runs estimates a run-to-run noise floor of ∼0.0025 nats, not a placement effect.The configurations are identical apart from data-order noise.Source: G MTP placement and weighting ablation · p. 12

H Feature-space contrastive auxiliaries are shortcut-prone

Feature-space contrastive auxiliaries can exploit position and document identity rather than future content, producing low auxiliary losses without useful future prediction. The observed shortcuts motivate, but do not resolve, the challenge of designing negatives that control both signals.

H Feature-space contrastive auxiliaries are shortcut-prone

  • With same-position, cross-sequence negatives, hybrid NTP+CPC reached ∼0.5 InfoNCE nats by combining position with document fingerprints, without requiring future-specific content.At 125M, K=2, its gap reached +0.154 nats, versus +0.028 for Transformer-aux and +0.017 for MTP at the same cut.Source: H Feature-space contrastive… · p. 12
  • Position narrows the candidates to a target column and document fingerprints identify the sequence, so the hybrid shortcut taxes the forwarded representation instead of requiring future-content prediction.The cross-entropy anchor prevented collapse, but did not prevent the shortcut-directed pressure.Source: H Feature-space contrastive… · pp. 12–13
  • At 125M, K=2, pure CPC collapsed into a position encoder; its stopped-run InfoNCE was ∼0.28 nats while the final head never learned.The run stopped at step 4,000 of 5,100; its auxiliary loss was far below the same-sequence chance level.Source: Table 8 · p. 12, H Feature-space contrastive… · p. 12
  • Both failures exploit identity signals that the negative set does not control; suppressing position and document identity together remains a possible but untested design direction.Designing such a negative-sampling scheme is beyond the study’s scope.Source: H Feature-space contrastive… · pp. 12–13

Related papers

Loading 2610.04867v1…