PaperlayerView exact revision on arXiv

Source-linked AI summary

Forward Pass Domain Adaptation (Without Cross-Layer Backpropagation)

Rivaan Patil, Simon Dennis, Hao Guo, Kevin Shabahang

arXiv:2608.14563v1cs.LGcs.AI

TL;DR

Fine-tuning large language models requires substantial memory for backward-pass activations, limiting adaptation in constrained or high-throughput settings. FPO uses late-layer output-error alignment and forward-pass-only updates, preserving off-domain benchmark performance within seed-noise of baseline across three model families.

  • Problem

    Fine-tuning large language models retains activations for backpropagation, creating a memory barrier that parameter-efficient methods do not eliminate.

  • Method

    FPO uses a layer-viability diagnostic and applies an output-derived pseudogradient to selected MLP layers without an autograd graph or backward pass through the model body.

  • Results

    Across three model families, FPO leaves off-domain benchmarks within seed-noise of baseline while full-network SFT and LoRA do not reproduce this behavior.

  • Takeaways & Limitations

    2.7–3.2× faster than full-network SFT and ∼40% lower in peak training memory, FPO supports late-layer adaptation in constrained training regimes.

  • Takeaways & Limitations

    FPO is evaluated for continued-pretraining next-token cross-entropy on three model families and has not been validated on mixture-of-experts, encoder-decoder, or non-causal architectures.

Abstract

from arXiv · show

Forward-Pass-Only MLP training (FPO) adapts large language models without a backward pass through the model body, achieving 2.7--3.2x the throughput of standard fine-tuning at ~40% less peak training memory, while leaving off-domain benchmarks within seed-noise of baseline, a property that full-network fine-tuning does not reliably reproduce. FPO rests on a single empirical observation: at late layers of a transformer, the output-layer prediction error approximates the true gradient with cosine similarity 0.47--0.59 across six public models we survey. We introduce a two-minute diagnostic that quantifies this approximation per layer for any model, identifying where late-layer adaptation is viable. Informed by the diagnostic, FPO computes a single error signal at the output and applies it to each target layer. No signal is propagated between layers, and no autograd graph is constructed at any point. We evaluate FPO on three model families (OLMo-2-7B, Qwen3-8B, Falcon3-7B). Across all three, FPO produces in-domain perplexity improvement and leaves MMLU, ARC-Challenge, HellaSwag, and Winogrande within seed-noise of baseline. Localizing SFT to FPO's target layers to enter this regime is also feasible, but at 2.2x the wall-clock cost of FPO.

1 Introduction

The paper introduces FPO, which adapts language models without constructing an autograd graph or backpropagating through the model body. It combines a two-minute layer-viability diagnostic with benchmark-safe late-layer adaptation, delivering higher throughput and lower memory use than SFT alternatives.

  • Motivation: 2–3× inference memory is typically required for fine-tuning because backward-pass gradient graphs store activations at every layer.This overhead is identified as the dominant barrier to fine-tuning on consumer GPUs, edge devices, and high-throughput personalization.
  • Motivation: 0.47–0.59 cosine similarity links output-projected pseudogradients to true MLP gradients in the last 25% of layers across six public transformers.For down-projection-specific gradients, the reported range is 0.35–0.59.
  • Contributions: Approximately two minutes on one GPU identifies viable target layers for any autoregressive transformer by measuring pseudogradient–gradient cosine similarity.The diagnostic is applied across six public models ranging from 3B to 8B parameters.
  • Contributions: 2.7–3.2× faster than full-network SFT, FPO adapts models without an autograd graph or backward pass through the model body.FPO is also 2.2× faster than SFT-partial and reduces peak training memory by approximately 40% at fixed batch size.
  • Contributions: Across OLMo-2-7B, Qwen3-8B, and Falcon3-7B, FPO keeps off-domain benchmarks within seed-noise of baseline while full-network SFT and LoRA do not.SFT-partial confirms that preservation is a regime property, while FPO is the practical method for operating in that regime.

2 Late-layer gradient redundancy and the cosine diagnostic

Late-layer output-projected pseudogradients approximate true gradient directions because they retain output-derived contributions while discarding later nonlinearities and cross-token routing. A cosine diagnostic identifies viable layers: across six transformers, late-layer cosines reach 0.47–0.59 and rise within the final quarter, while the diagnostic directly measures the cleanest W3 case.

  • Gradient decomposition: The true gradient decomposes into an output-computable pseudogradient plus a residual capturing discarded later-layer effects.The residual includes intermediate nonlinearities and attention’s cross-token routing through later layers.
  • Gradient decomposition: The pseudogradient combines the output residual-stream gradient with the cached forward activation feeding each MLP weight, using closed-form unembedding and final-RMSNorm Jacobians.These quantities are available from the forward pass or analytically from logits, without autograd machinery.
  • Cosine diagnostic: A roughly two-minute calibration on 16 sequences of 512 tokens measures cosine alignment between the pseudogradient and true gradient for each layer.Higher cosine means the pseudogradient captures the dominant true-gradient direction; near-zero cosine indicates it is uninformative.
  • Diagnostic scope and limitations: The diagnostic directly measures the W3 down-projection because its gradient is expressible from the layer-output gradient and cached MLP activation, whereas W1 and W2 require additional feedback approximations.FPO uses cached W3 feedback for W1 and W2, and those updates are validated empirically by downstream metrics.
  • Cosine diagnostic: Across six public 3B–8B transformers, late-layer cosines reach 0.47–0.59, with W3-specific cosine spanning 0.35–0.59 and the rise consistently occurring within the final quarter.Cosines are low and noisy through most layers before rising sharply and plateauing in late layers.

3 Forward-Pass-Only MLP Training (FPO)

FPO derives target-layer weight updates from one forward pass, using the analytically computed output error and closed-form normalization Jacobians without propagating gradients through the layer stack or constructing an autograd graph. Its stability depends on architecture-specific normalization corrections and a uniform RMS floor clamp.

  • Core method: FPO computes every target-layer update from the output-layer error and closed-form Jacobians, without any gradient signal traversing the layer stack.The output error is computed analytically from logits and applied directly to target-layer weights.
  • Training step: Each step caches target-layer MLP activations, forms e = softmax(z) − 1y, and projects it into residual-stream space with the transposed language-model head.The projected signal is the exact cross-entropy gradient with respect to the residual stream.
  • Training step: FPO uses the initialization feedback matrix for gate and up updates, which is exact at step 0 and becomes an approximation as W3 drifts.Late-layer adaptation and relative gradient clipping at 1% bound cumulative W3 drift over the operating range.
  • Memory: No autograd graph is constructed; peak memory consists of cached target-layer activations, optimizer state for target weights, and their float32 master copies.The target weights comprise ∼10% of parameters, keeping these costs small relative to an SFT gradient graph.
  • Normalization corrections: Correct normalization Jacobians and a uniform RMS floor of 0.5 are necessary for stability because otherwise per-token pseudogradient magnitudes become unstable.Pre-MLP Jacobian composition is applied for Qwen3 and Gemma 3 but not for Llama-family, Mistral, OLMo-2, or Falcon3.

4 Experimental setup

The experiments compare FPO with full-network SFT, LoRA, and target-layer SFT across three model families, using model-specific late-layer targets selected from cosine-curve plateaus. Adaptation uses continued-pretraining-style scientific or mathematical/biomedical corpora and evaluates in-domain perplexity alongside standardized off-domain benchmarks.

  • Models and target layers: FPO targets layers 27–31 in OLMo-2-1124-7B, 28–31 in Qwen3-8B, and 21–27 in Falcon3-7B-Base, selected from each model’s high-plateau cosine region.The models contain 32, 36, and 28 transformer blocks, respectively.
  • Adaptation domains: OLMo and Falcon3 use a 50/50 interleaved mix of mathematical and biomedical text, whereas Qwen3 uses mixed-subject scientific text from arXiv.The corpora are continued-pretraining-style, consisting of long-form documents tokenized in each model’s native vocabulary.
  • Methods compared: The comparison includes FPO with SGD and AdamW, full-network SFT, LoRA-16, and target-layer SFT on OLMo and Falcon3.LoRA-16 uses rank 16, α = 32, and all attention and MLP projections at every layer.
  • Evaluation: Evaluation combines in-domain perplexity with MMLU, ARC-Challenge, HellaSwag, Winogrande, and GSM8K on OLMo math.The benchmark settings are 5-shot, 25-shot, 10-shot, 5-shot, and additional OLMo GSM8K, respectively; OLMo’s baseline average is 61.49.
  • Compute and learning rates: All main results run on a single H100 80GB in bf16, while diagnostics and initial sweeps use an A100 40GB.Learning rates are selected through three-rate sweeps for SFT, LoRA, and FPO, with full sweeps reported in Appendix B.

5 Results

FPO delivers 2.7–3.2× higher throughput than full SFT while using substantially less memory, and late-layer adaptation preserves off-domain benchmarks across model families. This preservation comes with smaller in-domain perplexity gains than methods that update earlier layers.

  • Throughput and memory: 2.7–3.2× faster than SFT-full across three model families, with larger feasible batch sizes on identical hardware.FPO eliminates the gradient graph; on OLMo-2-7B at sequence length 2048, it uses 17.7 GB at bs=1 versus 29.6 GB for SFT.
  • Throughput and memory: 2.2× throughput advantage over SFT-partial on Falcon3: FPO reaches 22.7K tok/s at bs=10 versus 10.1K tok/s.SFT-partial remains bottlenecked by backward computation through trainable layers and activation storage across the layer stack.
  • Off-domain benchmark preservation: FPO keeps off-domain benchmark averages within 0.22 points of baseline on Qwen3 and Falcon, while improving OLMo by +0.91.On OLMo, this pattern replicates with SFT-partial at +0.94, whereas full-network methods do not consistently preserve benchmarks.
  • Off-domain benchmark preservation: σ = 0.02 for FPO’s OLMo benchmark average across three data orderings, with per-benchmark σ ≤0.09.The reported per-benchmark values are MMLU 0.02, ARC-C 0.05, HellaSwag 0.03, and Winogrande 0.09.
  • Layer-selection mechanism: Lower diagnostic cosine values correspond to narrower preservation advantages over full-network methods; Qwen3 has the smallest cost and narrowest FPO advantage.OLMo and Falcon have higher cosines, larger full-network benchmark damage, and larger FPO preservation advantages.
  • Perplexity trade-off: FPO’s in-domain perplexity improvement is smaller than SFT-full’s or LoRA’s in every setting, while held-out text perplexity stays within 0.5% of baseline.LoRA instead degrades held-out text perplexity by +4.1% and +19.4% on two of three models.

6 Discussion and limitations

FPO is most useful when memory, benchmark stability, or rapid iteration matter, while SFT-full remains appropriate for maximizing domain perplexity without regard to off-domain behavior. Its forgetting resistance arises from gradient decomposition, but its current scope is next-token continued-pretraining adaptation and excludes several architectures and objectives.

  • When should practitioners use FPO?: FPO suits memory-constrained, benchmark-sensitive, and rapid-development settings, whereas SFT-full targets maximum domain perplexity without off-domain constraints.The passage identifies consumer GPUs, edge devices, batched personalization, off-domain stability, and throughput across iteration cycles as motivating regimes.
  • Relationship to catastrophic forgetting: FPO’s benchmark preservation is forgetting resistance derived from gradient decomposition rather than regularization or a continual-learning protocol.The passage expects this resistance to compose naturally with existing continual-learning approaches.
  • Architectural and methodological scope: FPO currently targets continued-pretraining adaptation with next-token cross-entropy and has not been validated on mixture-of-experts, encoder-decoder, or non-causal architectures.Extending FPO to instruction-tuning, RL-style objectives, or preference optimization requires re-deriving the output-projected pseudogradient; the diagnostic covers six public causal-LM architectures, while evaluation spans three families and two normalization layouts.

7 Related work

FPO differs from prior forward-only, zeroth-order, parameter-efficient, memory-efficient, selective-layer, and forgetting-resistance methods in both mechanism and justification. It eliminates the model-body backward pass while using layer selection and update behavior to preserve transfer.

  • Forward-only and feedback-alignment methods: FPO and direct feedback alignment both use fixed feedback multiplication, but DFA relies on feedback unrelated to the Jacobian whereas FPO assumes initialized feedback remains close under bounded weight movement.The methods share a multiplication but are justified in opposite directions.
  • Zeroth-order methods: FPO uses exact gradients on selected parameters, unlike MeZO’s finite-difference estimates, whose high variance and slow convergence require orders of magnitude more steps for a given loss.FPO’s per-step behavior is comparable to standard SFT on those parameters.
  • Parameter-efficient fine-tuning: 2.4× higher throughput: FPO outperforms LoRA on OLMo-2-7B at batch size 1, while LoRA’s peak memory is within 0.4 GB of FPO’s.Both methods reduce memory differently: LoRA reduces optimizer-state costs, whereas FPO avoids the backward pass.
  • Memory-efficient training: FPO is orthogonal to quantization and low-rank gradient projection because it eliminates the model-body backward pass entirely rather than reducing the cost of constructing the full gradient graph.QLoRA quantizes weights, while GaLore projects gradients into smaller subspaces.
  • Selective-layer methods: FPO extends selective-layer methods with a diagnostic for where backpropagation can be skipped and an update that dispenses with autograd machinery.Prior layer freezing and surgical fine-tuning restrict updates to particular layers to improve transfer or reduce forgetting.
  • Catastrophic forgetting: For single-domain adaptation, restricting updates to late layers provides forgetting resistance without explicit regularization or replay, because the mechanism follows from update location.This connects FPO’s benchmark-preservation property to catastrophic-forgetting research while making a more targeted claim.

8 Conclusion

FPO exploits the observation that late-layer gradients are closely aligned with a forward-pass-computable output quantity. Across three model families, this enables late-layer adaptation in memory regimes where backpropagation is infeasible.

  • 8 Conclusion: Cosine 0.47–0.59: late-layer true gradients align with a quantity computable from the forward pass alone.This alignment provides the empirical basis for FPO.
  • 8 Conclusion: FPO turns the 2–3× memory overhead and substantial throughput cost of training into negligible additions over inference.The method operationalizes the late-layer gradient-alignment observation without requiring conventional training overhead.
  • 8 Conclusion: Across three model families, FPO enables late-layer adaptation in memory regimes where backpropagation is infeasible.The conclusion reports this feasibility across the evaluated model families.

A Implementation details … C.2 Cross-model implications for learning rate

FPO uses cached forward-pass signals, norm corrections, clipped updates, and float32 master weights to adapt target-layer MLPs without cross-layer backpropagation. Learning-rate sweeps and norm diagnostics show that FPO optima vary substantially across models, while SFT and LoRA remain insensitive within stable ranges.

  • A Implementation details: FPO projects output-logit error into residual-stream space, applies final- and post-FFN-norm Jacobian corrections, and updates target-layer W3 from cached activations.The W3 update uses dW3 = einsum('bld,blg->dg', g_ell, h_mid) / (B * L), followed by relative-norm clipping.
  • A Implementation details: Float32 master weights prevent ∼10−6 updates from rounding to zero when bf16 represents ∼10−4 changes for weights near 10−1.Updates are applied to float32 master copies and cast to bf16 for forward passes.
  • B Full LR sweeps: 50M-token learning-rate sweeps use batch size 8 for FPO and batch size 1–2 for SFT/LoRA, with sequence length 2048 and bf16 on one H100 80GB.These settings define the appendix comparisons across model families.
  • B.1 OLMo-2-7B (math + bio, 50M tokens, target layers 27–31): SFT-full and LoRA-16 converged identically across their swept ranges on OLMo-2-7B, yielding math −11.4% and −10.2%, respectively.The corresponding bio results were −8.3% for SFT-full and −7.7% for LoRA-16.
  • B.2 Qwen3-8B (scientific, 50M tokens, target layers 28–31): 1 × 10−2 is Qwen3-8B’s FPO optimum, with degradation at ≥2 × 10−2 and a 20-fold gap versus OLMo’s 5 × 10−4 optimum.The ordering is qualitatively consistent with per-layer update-magnitude ratios.
  • B.3 Falcon3-7B (math + bio, 50M tokens, target layers 21–27): SGD 0.003 is Falcon3-7B’s best FPO configuration, producing math −5.9%, bio −2.3%, text +0.4%, and code +1.3%.The SGD sweep included 0.003 and 0.01; AdamW was also swept over 5e-4, 1e-3, and 3e-3.
  • B.4 Summary: Across all three models, SFT and LoRA are learning-rate-insensitive within swept stable ranges, whereas FPO optima are 5e-4, 1e-2, and 1e-3.Runs above the stable ranges degraded or diverged, while runs below them were undertrained at 50M tokens.
  • C Norm diagnostics and pseudogradient calibration / C.1 OLMo-2-7B per-layer measurements: The norm diagnostic calibrates FPO learning rates because pseudogradient magnitude and target-weight norms determine the fractional weight change produced by a given learning rate.For OLMo-2-7B, the min-RMS clamp at 0.5 prevents outlier-token amplification of up to 12×, while the implied SGD LR matches the empirical optimum of 0.003.

C.3 Ablations … D.2 Qwen3-8B extended trajectory (FPO AdamW, lr 1e-2, 0–500M tokens)

The ablations identify robust target-layer, optimizer, and learning-rate choices, while extended trajectories show domain gains can persist but cached-Jacobian stability differs across models. OLMo remains comparatively stable off-domain, whereas Qwen3’s long-horizon text perplexity reverses upward as W3 drifts.

  • C.3 Ablations: −5.7% to −5.8% math improvement is essentially unchanged across 3, 5, or 8 target layers, while text perturbation increases only +0.1% to +0.3%.Adding lower-cosine layers neither improves domain gain nor substantially harms it; 3–5 top layers are recommended for 28–32-layer models.
  • C.3 Ablations: Math endpoints differ by ∆within 1 point between FPO-AdamW and FPO-SGD on OLMo.AdamW is preferred when memory permits, while SGD suits the most memory-constrained settings.
  • C.3 Ablations: A roughly 5× learning-rate window around the optimum preserves FPO’s peak quality across SGD 0.001–0.03 and AdamW 1e-5 to 1e-2 sweeps.The result indicates that practitioners need not sweep learning rates finely.
  • D Extended trajectories and cached-Jacobian stability: Longer-horizon trajectories and W3-drift measurements support the cached-Jacobian assumption over the operating range.These measurements are reported to assess extended training behavior and cached-Jacobian stability.
  • D.1 OLMo-2-7B extended trajectory (FPO AdamW, lr 5e-4, 0–500M tokens): −6.3% math perplexity is reached by ∼50–100M tokens on OLMo and remains near saturation through 500M.Bio improves from −0.7% at 50M to −1.9% at 500M, while text ends at +1.0% after slowly drifting.
  • D.1 OLMo-2-7B extended trajectory (FPO AdamW, lr 5e-4, 0–500M tokens): +1.0% text perplexity drift after 500M tokens remains below SFT-full’s +1.8% and LoRA’s +4.1% at their 50M comparison budget.FPO was trained 10× longer than the comparison budget while showing less off-domain disturbance.
  • D.2 Qwen3-8B extended trajectory (FPO AdamW, lr 1e-2, 0–500M tokens): −11.0% scientific perplexity improvement at 500M tokens accompanies text perplexity reversing from −4.2% to +12.3% on Qwen3.The profile is consistent with extended training pushing W3 farther from the cached W3 init used in W1/W2 updates; absolute drift reaches a larger fraction of the initial weight norm.

D.3 W3 drift over training · E Architecture-specific Jacobian derivations · E.1 RMSNorm Jacobian

The cached-Jacobian assumption remains empirically supported over the 50M operating range, while architecture-specific FPO Jacobians are composed from output errors and intervening RMSNorm Jacobians computed directly from cached activations. A minimum-RMS clamp limits amplification from anomalously low-RMS tokens.

  • D.3 W3 drift over training: ≤0.6% fractional W3 drift for OLMo and Falcon at 50M accompanies approximately stable text perplexity.The absolute Frobenius drift is ≤0.7 against initial weight norms of ∼120–127.
  • D.3 W3 drift over training: ∼1.2% fractional W3 drift for Qwen3 at 50M grows to ∼5% at 500M, coinciding with text-perplexity reversal.This drift occurs under Qwen3’s aggressive learning rate.
  • D.3 W3 drift over training: The cached-Jacobian assumption is empirically supported over the 50M operating range across OLMo, Falcon, and Qwen3.The passage links this support to the observed drift and perplexity behavior across the three models.
  • E Architecture-specific Jacobian derivations: FPO composes the output-layer error with the Jacobians of each RMSNorm between the output and target layer.The appendix provides architecture-specific derivations and compositions.
  • E.1 RMSNorm Jacobian: The RMSNorm Jacobian-vector product is derived for a D-dimensional vector x using mean(x2) + ε and upstream gradient g.This supplies the RMSNorm component used in the architecture-specific composition.
  • E.1 RMSNorm Jacobian: RMSNorm Jacobian expressions are computed directly on cached forward-pass activations, with no autograd graph constructed at any point.The computation uses ˆx = x/ρ alongside mean(x2) + ε.
  • E.1 RMSNorm Jacobian: ρmin = 0.5 caps 1/ρ amplification for outlier tokens with anomalously low post-norm RMS, while remaining inactive on typical tokens.The example given is the BOS token at certain layers in OLMo-2.

E.2 Per-architecture composition · F Evaluation harness and reproducibility

Architecture-specific FPO composition selects normalization Jacobians according to each model family’s layout, while evaluation uses standardized lm-eval-harness configurations and documented implementation details for reproducibility.

  • E.2 Per-architecture composition: The output-layer error is projected through lm_head transpose into residual-stream space and always composed with the final RMSNorm Jacobian.This composition applies across all surveyed architectures.
  • E.2 Per-architecture composition: Post-feedforward RMSNorm composition is used for OLMo-2 and Gemma 3, with γpffn values of 0.25–0.39 in OLMo target layers.The min-RMS clamp prevents destabilization on outlier tokens.
  • E.2 Per-architecture composition: Pre-MLP RMSNorm composition is used for Qwen3 and Gemma 3, but omitted for Llama-family, Mistral, OLMo-2, and Falcon3.The decision is fixed once per model family because it was neutral or mildly destabilizing in the omitted families.
  • E.2 Per-architecture composition: Each new architecture family requires identifying its normalization layout once; the same closed-form Jacobian is then applied at every target location.The resulting family configurations are Llama/Mistral/Falcon3: final norm only; OLMo-2: final plus post-FFN; Qwen2/Qwen3: final plus pre-MLP; Gemma 3: all three.
  • F Evaluation harness and reproducibility: All Section 5.2 benchmark numbers use lm-eval-harness with MMLU, ARC-Challenge, HellaSwag, Winogrande, and OLMo-only GSM8K configurations.Evaluations run in bf16, with GSM8K reported using flex-extract and strict-match.
  • F Evaluation harness and reproducibility: Evaluation uses each model’s native tokenizer without substitution and runs benchmarks on a single H100 80GB or A100 80GB.Perplexity-only evaluations use a single A100 40GB.
  • F Evaluation harness and reproducibility: The method, architecture-specific Jacobian compositions, per-layer norm measurements, and complete learning-rate sweeps provide sufficient detail for independent reimplementation.Implementation and evaluation configurations will be released alongside the camera-ready version.
Loading 2608.14563v1…