Which Layers Need Backpropagation? Position-Dependent Partial Training for Memory-Efficient LLM Pretraining

Read full text
Authors: Raen2264
Cited by: 2

Abstract

We ask which layers of a transformer actually need backpropagation during pretraining from scratch. Training Llama-1B with backpropagation restricted to a sliding window of layers (plus embeddings and the LM head) and freezing the remainder at random initialization, we find that the best training method depends critically on the training horizon. budget and schedule, frozen-partial reaches the lowest perplexity (35.2), beating full backpropagation (39.0) by 11% and GaLore (52.4) by 33%, while using 37% less VRAM at a matched micro-batch (10.9 vs 17.4 GB). A fine-grained sweep locates the crossover between 25M and 50M tokens on a single dataset, ruling out a dataset confound. A position sweep reveals a U-shaped curve whose optimum lies in the early-middle layers, coinciding with the network's lowest gradient signal-to-noise region. The non-BP layers can be frozen entirely: zeroth-order perturbation contributes at most ~4 PPL beyond freezing. Frozen-partial is also the most robust to learning-rate choice. The optimal BP position is architecture-dependent: attention models prefer an early-middle window, while recurrent/SSM models require the input side. All results are at the 1B scale; cross-scale (3B/7B) study is left to future work. Companion code and experiment logs: https://github.com/2264K/hybrid-zo-pretrain License: paper CC BY 4.0; code Apache-2.0.

Fields of Study & Subjects

Advanced Neural Network Applications Advanced Memory and Neural Computing Machine Learning in Materials Science

Publication details

Journal / source
Zenodo (CERN European Organization for Nuclear Research)
Publication date
26 May 2026
Language
EN
Document type
Preprint