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

AI summary

Overview

• Research area: Machine learning — decentralized and model-parallel training of language models, specifically greedy block-wise local learning as an alternative to end-to-end backpropagation.
• Technical level: Intermediate. The paper uses standard transformer training vocabulary (blocks, auxiliary heads, next-token prediction, validation NLL), but the core ideas — gradient isolation, block boundaries, per-block memory — are explained well enough for readers familiar with basic LM training.
• One-sentence scope: A token-budget-matched empirical study at 125M and 400M parameters that measures the held-out likelihood gap between greedy local learning and end-to-end backpropagation across block counts and auxiliary designs, and reports the memory consequence of per-block execution.

What This Paper Is About

Greedy local learning splits a network into gradient-isolated blocks, each trained by its own auxiliary loss, so that no error signal crosses block boundaries. This makes inter-stage communication forward-only and lets every block update its optimizer independently — exactly the properties wanted for model-parallel training over weakly connected or memory-constrained workers.

Prior work showed local learning is competitive with backpropagation on image classification, and that on small transformers it trades a worse best loss for parallel speedup. What had not been measured is how that likelihood gap behaves in autoregressive LM pretraining at larger scale, and which auxiliary design choices reduce it. This paper fills that gap with a controlled study at 125M and 400M parameters under Chinchilla-optimal token budgets.

Key Contributions

1. A token-budget-matched likelihood study of greedy local learning for LM pretraining at 125M and 400M parameters, across K ∈ {1, 2, 4} equal-depth blocks and Chinchilla-optimal budgets, extending the small-transformer language-modeling results of Laskin et al. (2021) to larger scale.
2. Per-block diagnostic evidence consistent with what the authors call representation myopia: later blocks barely improve on the representation they inherit under local NTP supervision.
3. A factorized auxiliary-design comparison showing that a Transformer-based auxiliary is a strong network-side intervention, while multi-token prediction (MTP) helps at the first block boundary but hurts when added at deeper boundaries — together with a reported failure of an NTP+CPC contrastive hybrid.
4. Deployment-style memory measurements and analytical accounting of inter-stage communication and auxiliary training compute for local learning as decentralized model parallelism, plus release of the factorized implementations used in the comparisons.

Main Findings

• The isolation gap grows with block count. For every auxiliary design, the validation NLL gap versus end-to-end training more than doubles from K=2 to K=4 (for example, MLP-aux at 125M: +0.046 → +0.132 nats). This matches the direction reported by Laskin et al. (2021) for small transformers.
• The gap shrinks with scale at K=4. Moving from 125M to 400M reduces the K=4 gap for all three designs: MLP-aux +0.132 → +0.112, Transformer-aux +0.078 → +0.075, and first-block MTP +0.074 → +0.062. This trend is specific to the deeper partition — at K=2, two of the three gaps increase slightly with scale.
• The lowest measured isolation cost is +0.027 nats. At the best operating point (first-block MTP, K=2, 400M), the cost of gradient isolation is +0.027 nats, about 1.1% relative validation NLL. End-to-end baselines are 2.771 (125M) and 2.418 (400M) nats/token.
• Local losses exhibit representation myopia. At 125M, K=4, the four per-block local loss curves collapse onto nearly the same trajectory under both NTP auxiliaries: the first block already reaches almost the loss of the full stack, and blocks above it barely improve on what they inherit. Each block behaves as a standalone next-token predictor rather than a feature extractor for its successors. The same qualitative pattern holds at 400M.
• A Transformer-based auxiliary helps most. Replacing MLP-aux with Transformer-aux reduces the gap by 39% and 41% at 125M (K=2 and K=4) and by 22% and 33% at 400M. The per-block curves move downward without separating, so the stronger auxiliary makes locally predictive features easier to decode but does not remove the myopic pattern. The authors note the two heads are not parameter- or compute-matched, so this is an architecture-level comparison rather than a controlled causal estimate of attention alone.
• MTP depends on boundary position. At K=2 the single boundary is necessarily the first, and MTP improves over Transformer-aux by 0.011 nats at 125M and 0.005 at 400M. At K=4, first-block MTP improves over Transformer-aux (+0.074 vs +0.078 at 125M; +0.062 vs +0.075 at 400M), whereas all-block MTP is worse than Transformer-aux at both scales (+0.082 and +0.086). The 400M first-block margin over Transformer-aux is 0.013 nats, about five times the ~0.0025-nat run-to-run noise floor; the 125M margin (0.004) is only consistent in direction.
• A per-offset diagnostic explains the MTP trade-off. In all-block MTP at 125M, K=4, the second-token weight decreases across successive auxiliary blocks (0.75, 0.50, 0.25), and next-token cross-entropy improves from 3.087 to 2.858, while second-token CE stays around 4.5–4.6. Lower blocks give up immediate predictive sharpness under the stronger forward-looking pressure. The authors hypothesize that early MTP preserves broad future-predictive information that many remaining layers can exploit, whereas near the final head the more uncertain second-token target competes with NTP with too little downstream depth to repay the loss.
• Placement matters more than schedule. Moving MTP pressure from all boundaries to the first improves validation NLL by 0.008 nats (125M) and 0.024 nats (400M) at K=4, while changing the weighting within all-block MTP (linear versus uniform) changes results by ~0.003 nats, at the level of noise. The all-block uniform configuration was not run at 400M.
• Feature-space contrastive auxiliaries were shortcut-prone. A pure CPC-style InfoNCE auxiliary (125M, K=2, same-sequence negatives) collapsed into a content-free position encoder: auxiliary InfoNCE reached ~0.09 nats at its lowest and ~0.28 at the stopping checkpoint, while the final-head loss stalled at ~7.25 nats against ~2.9 for NTP-trained auxiliaries, and the run was stopped at step 4,000 of 5,100. After switching to same-position, cross-sequence negatives, an NTP+CPC hybrid still drove InfoNCE to ~0.5 nats — below the log B ≈ 2.8 bound for batch size B — by combining position with coarse document fingerprints. That hybrid increased the gap to +0.154 nats at a single cut, worse than MLP-aux with three cuts (+0.132), versus +0.028 for Transformer-aux and +0.017 for MTP at the same cut.
• Local execution cuts memory. At 400M, splitting into K=4 blocks cuts peak per-worker GPU memory from 29.2 to 14.5–15.8 GiB (~1.9–2.0×) and activation memory from 27.7 to 12.4–13.4 GiB (up to 2.2×), depending on the auxiliary head. The reduction is sublinear in K because every worker still runs a vocabulary-sized auxiliary head, while the transformer-block activations themselves scale as 1/K. First-block MTP carries the MTP head only at the first stage, and its maximum-stage footprint at K=4 is within 0.1 GiB of plain Transformer-aux.
• Memory savings come with training compute overhead. One auxiliary head costs 19%, 26%, and 45% of end-to-end training FLOPs at 125M for MLP-aux, Transformer-aux, and an MTP head, and 9%, 13%, and 22% at 400M. At K=4 the totals are 78% for Transformer-aux and 97% for first-block MTP at 125M, and 39% and 48% at 400M. Most of this is the vocabulary projection every local head needs — 16.0% of end-to-end FLOPs per head at 125M and 7.1% at 400M — which an MTP head runs twice. Compute alone does not order the results: all-block MTP costs the most and is worse than Transformer-aux at K=4.
• Inter-stage traffic halves analytically. Pipeline-parallel backpropagation ships activations forward and gradients backward (2bd total per token per boundary) and is update-locked; local learning deletes the backward direction (bd per token per boundary) with no update locking. For the 400M class (d=1,024, bf16) this is 2 KiB versus 4 KiB per token per boundary. MTP targets add no inter-stage traffic because they are already in the token stream every worker holds.

Methodology in Plain English

The researchers took Llama-style decoder-only transformers at two sizes — a 125M-class model (d_model=768, 12 layers, 12 heads, 85M non-embedding parameters) and a 400M-class model (d_model=1024, 24 layers, 16 heads, 308M non-embedding parameters) — and split each into K ∈ {1, 2, 4} contiguous, equal-depth blocks. They trained on the FineWeb-Edu 10BT sample packed into 4,096-token sequences using the Llama-2 32k SentencePiece tokenizer, at Chinchilla-optimal budgets of 20 tokens per total parameter: 5,100 steps (≈2.7B tokens) for the 125M class and 14,300 steps (≈7.5B tokens) for the 400M class, with a global batch of 2^19 (≈524k) tokens per step.

Gradient isolation works by stopping gradients at each block boundary: activation values are identical to the original network, but each block and its auxiliary head receive gradient only from their own local loss. The K=1 configuration is bit-identical to standard end-to-end training and serves as the baseline.

The design study factorizes auxiliary design into two choices. On the network side, they compare an MLP head (a pre-normalized SwiGLU FFN with residual, RMSNorm, and a vocabulary projection) against a Transformer head (one additional transformer block prepended to the same head). On the objective side, they compare standard next-token prediction against multi-token prediction with two offset branches (j=1 predicting x_{t+1}, j=2 predicting x_{t+2}) sharing the final RMSNorm and vocabulary projection, in either a first-block-only or an all-block configuration. Heads are untied across blocks, trained only through their local losses, and discarded at inference.

Evaluation uses held-out validation NLL per token on FineWeb-Edu, always computed through the final block's real NTP head so it is directly comparable across methods, block counts, and to end-to-end training. Optimizer settings are shared across methods and block counts within each scale, all quality experiments are synchronous simulations with the same seed, and the memory benchmark executes each stage separately on a single GPU at batch 4×4,096 tokens in bf16.

Why This Matters

Impact on research. The paper moves local learning from image classification and small transformers into the regime researchers actually care about for pretraining — hundreds of millions of parameters under standard Chinchilla token budgets, judged by held-out likelihood rather than wall-clock proxies. Its most distinctive contributions are negative or diagnostic rather than a new winning method: the observation that later blocks barely improve the representations they inherit, the finding that the same auxiliary objective helps at one boundary and hurts at another, and the documented failure of contrastive auxiliaries that get satisfied by position and document identity rather than future content. That last result is a concrete warning for anyone porting self-supervised objectives into local learning.

Real-world applications the paper's setting motivates:

• Pretraining or continued pretraining

Authors’ abstract

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.