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

2026-10-04

AI Summary

This study addresses the significant loss gap between greedy chunk-wise local learning and end-to-end training in large model pre-training. Operating within the 125M to 400M parameter regime under Chinchilla-optimal token budgets, this work proposes Transformer auxiliary networks combined with a selective multi-token prediction (MTP) objective, integrated with decentralized model parallelism, to optimize training at deep block boundaries. The findings reveal that with K=4, the performance gap narrows as model scale increases, and the proposed auxiliary networks recover 22%–41% of this discrepancy. Notably, applying MTP exclusively to the first chunk yields optimal results. Furthermore, the approach achieves a 2.2× reduction in activation memory, providing an effective strategy for efficient local training.

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.

Problem

• greedy local learning
• language model pretraining
• auxiliary loss design
• loss gap
• decentralized training

Innovation

• Greedy Local Learning
• Block-wise Training
• Auxiliary Objective Design
• Multi-Token Prediction
• Decentralized Model Parallelism