🤖 machine learning

Forward Pass Domain Adaptation (Without Cross-Layer Backpropagation)

This paper introduces Forward-Pass-Only (FPO) training, a domain adaptation method for large language models that achieves significantly higher throughput and lower memory usage by applying a single output-layer error signal directly to target layers without backpropagation, while maintaining performance comparable to standard fine-tuning on both in-domain and off-domain benchmarks.

Original authors: Rivaan Patil, Simon Dennis, Hao Guo, Kevin Shabahang

Published 2026-08-18

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Rivaan Patil, Simon Dennis, Hao Guo, Kevin Shabahang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models are the engines behind many of the most advanced artificial intelligence tools available today, capable of writing, reasoning, and solving problems with a fluency that once seemed impossible for machines. To make these models useful for specific tasks, such as diagnosing medical conditions or analyzing legal documents, researchers typically "fine-tune" them. This process involves showing the model vast amounts of new, specialized data and adjusting its internal settings so it learns the patterns of that new world. However, this adjustment is incredibly expensive. It requires powerful computers to run the model forward to see the answer, and then backward to figure out exactly how to change every single part of the system to get a better result. This backward step demands a massive amount of memory, often three times what is needed just to run the model, making it impossible to perform on standard consumer computers or in situations where speed is critical.

A team of researchers at the University of California, Santa Cruz, and the University of Melbourne has discovered a way to bypass this expensive backward step entirely. They found that for large language models, the most important adjustments needed to learn new information happen mostly in the final layers of the network, near the end of the processing chain. By focusing only on these late layers and calculating the necessary changes using only the information available from a single forward pass, they created a new training method called Forward-Pass-Only MLP training. This approach allows the model to learn new skills without ever looking backward through its own structure, reducing the memory required by about forty percent and making the process nearly three times faster than standard methods. Crucially, this speed and efficiency come without the usual cost of forgetting what the model already knew; while traditional methods often degrade the model's general abilities when learning a new specialty, this new technique keeps the model's original knowledge intact.

The core of this discovery rests on a simple but profound observation about how these models process information. In a standard training session, the computer calculates an error signal at the very end of the model's output and then sends that signal backward, layer by layer, to adjust the weights inside the network. The researchers realized that for the final quarter of the model's layers, the direction of the necessary adjustment is almost identical to a signal that can be calculated using only the output and the current state of the model, without needing to trace the path backward. They tested this idea across six different public models, ranging from three billion to eight billion parameters, and found that the direction of the true adjustment aligns closely with this simpler, forward-only signal. The alignment was strong enough to be useful, with the signals pointing in roughly the same direction in nearly half of the cases, a consistency that held true across different model architectures.

To put this theory into practice, the team developed a quick diagnostic tool that takes about two minutes to run on a single computer chip. This tool measures how well the simple forward signal matches the complex backward signal for each layer of a specific model. It acts as a map, showing researchers exactly where in the network the backward pass can be safely skipped. Once these "viable" late layers are identified, the training process changes fundamentally. Instead of building a massive, complex record of every step the model took so it can reverse them later, the system simply runs the model forward, calculates the error at the end, and applies a direct correction to the target layers. No backward pass is ever constructed, and no memory is wasted storing the history of the computation. This eliminates the primary bottleneck that has prevented many people from fine-tuning large models on their own hardware.

The results of testing this method on three different families of models were striking. In terms of speed, the new approach was between 2.7 and 3.2 times faster than standard fine-tuning, allowing researchers to process significantly more data in the same amount of time. On memory-constrained hardware, such as a single high-end graphics card, the method allowed for much larger batches of data to be processed simultaneously, whereas standard methods would crash due to memory limits. Perhaps most importantly, the method preserved the model's general intelligence. When the researchers tested the adapted models on a variety of unrelated tasks, such as answering general knowledge questions or solving logic puzzles, the models performed almost exactly as well as they did before the training. In contrast, models trained with standard methods often saw their performance on these general tasks drop significantly, a phenomenon known as catastrophic forgetting. The new method avoided this trap, keeping the model's broad capabilities stable while it learned its new specialty.

The researchers also explored why this preservation of general knowledge happens. They compared their method to a hybrid approach where they only updated the late layers but still used the standard, slow backward pass. They found that the preservation of general knowledge was not a trick of the new calculation method itself, but a direct result of restricting the changes to the late layers of the network. The early layers, which handle the fundamental building blocks of language and reasoning, were left untouched, preventing the model from overwriting its core understanding. However, the new method is the only practical way to achieve this, because the standard backward pass is too slow and memory-intensive to be feasible when restricted to just a few layers. The new approach makes it possible to operate in this "safe" zone where learning is efficient and forgetting is minimized.

While the method is highly effective, it does have limits. The models trained with this technique did not learn the new domain as deeply as models trained with the full, slow method; they achieved slightly less improvement on the specific task they were taught. This is the trade-off: the new method sacrifices a small amount of maximum performance on the new task to gain massive improvements in speed, memory efficiency, and the preservation of general knowledge. The researchers suggest that for many real-world applications, such as personalizing a model for a specific user or adapting it for a specialized industry on consumer hardware, this trade-off is well worth it. The method opens the door for adaptation in settings where it was previously impossible, turning a process that once required a data center into something that can run on a single machine, all while keeping the model's original intelligence safe from the damage of over-specialization.

Technical Summary: Forward Pass Domain Adaptation (FPO)

1. Problem Statement

Fine-tuning large language models (LLMs) on domain-specific data typically requires a full backward pass through the network. This process necessitates storing activations for every layer to compute gradients, resulting in a memory footprint 2–3 times larger than inference. This overhead prevents fine-tuning on consumer GPUs, edge devices, or in high-throughput personalization scenarios. While parameter-efficient methods like LoRA reduce optimizer state, they still require a full backward pass, thereby materializing the full activation graph. Furthermore, full-network fine-tuning often leads to "catastrophic forgetting," where performance on off-domain benchmarks (e.g., MMLU, ARC-Challenge) degrades significantly, a behavior that is not reliably reproducible across different random seeds.

2. Core Observation and Methodology

The paper posits that at the late layers of a transformer, the true gradient with respect to MLP weights can be approximated by a quantity computable solely from the forward pass.

2.1 Theoretical Basis

For an MLP weight $W^{(\ell)}$ at layer $\ell$ , the true gradient $\nabla W^{(\ell)}L$ decomposes into:
$\nabla W^{(\ell)}L = \tilde{g}W^{(\ell)} + R W^{(\ell)}$
Where:

$\tilde{g}W^{(\ell)}$ is the output-projected pseudogradient. It is computed by backpropagating the output-layer error through the unembedding transpose and the closed-form Jacobians of the final normalization layers only. It assumes all intermediate layers between $\ell$ and the output act as identity functions.
$R W^{(\ell)}$ is the residual, containing contributions from intermediate nonlinearities and cross-token routing.

Empirical surveys of six public models (3B–8B parameters) reveal that in the last 25% of layers, the cosine similarity between $\tilde{g}W^{(\ell)}$ and the true gradient ranges from 0.47 to 0.59. This suggests the residual is small enough in the orthogonal direction that the output-projected component captures the dominant gradient direction.

2.2 The Diagnostic Tool

The authors introduce a layer-viability diagnostic that quantifies this cosine similarity per layer.

Process: It runs a small calibration batch (16 sequences of 512 tokens) on a single GPU, taking approximately two minutes.
Output: It identifies the "high-plateau" region in the final layers where the alignment is strong enough to exploit. This allows practitioners to determine where in a specific model late-layer adaptation is viable.

2.3 Forward-Pass-Only MLP Training (FPO)

FPO is an algorithm that adapts language models without constructing an autograd graph or running a backward pass through the model body.

Mechanism:
1. Perform a standard forward pass, caching MLP intermediate activations ( $h_{mid}$ ) at target layers.
2. Compute the output-layer error ( $e = \text{softmax}(z) - y$ ).
3. Project $e$ into residual-stream space using the language model head transpose ( $g = e W_{lm}$ ).
4. Compose $g$ with closed-form Jacobians of normalization layers (e.g., RMSNorm) to produce the layer-specific signal $\tilde{g}^{(\ell)}$ .
5. Weight Updates:
  - Down-projection ( $W_3$ ): Updated exactly using $\Delta W_3 \propto -\tilde{g}^{(\ell)T} h_{mid}$ .
  - Gate ( $W_1$ ) and Up ( $W_2$ ) projections: These require a gradient signal in the MLP hidden space. FPO approximates this by using the cached initial weights ( $W_3^{init}$ ) as a fixed feedback matrix: $\text{error}_h = \tilde{g}^{(\ell)} W_3^{init}$ . This approximation holds because $W_3$ changes slowly under the adaptation regime.
Constraints: No signal is propagated between layers; no autograd graph is constructed. The peak memory overhead is limited to cached activations at target layers and optimizer states for those specific layers (approx. 10% of parameters).

3. Key Contributions

Layer-Viability Diagnostic: A fast, two-minute tool that quantifies the cosine similarity between output-projected pseudogridients and true gradients, identifying viable target layers for any auto-regressive transformer.
FPO Algorithm: A training method that achieves domain adaptation without a backward pass through the model body. It is 2.7–3.2× faster than full-network supervised fine-tuning (SFT) and 2.2× faster than partial-layer SFT (SFT-partial). It reduces peak training memory by ~40% compared to standard SFT.
Benchmark-Safe Regime: Empirical evidence showing that restricting adaptation to late layers (via FPO) preserves off-domain benchmark performance within seed-noise of the baseline. In contrast, full-network SFT and LoRA often degrade these benchmarks significantly.

4. Experimental Results

The method was evaluated on three model families: OLMo-2-7B, Qwen3-8B, and Falcon3-7B, adapted on mathematical, biomedical, and scientific text.

Throughput and Memory:
- FPO achieves 22.0K–22.7K tokens/s on H100 GPUs, compared to 7.5K–8.7K tokens/s for full SFT.
- FPO supports larger batch sizes (e.g., batch size 8–10) on hardware where SFT fails to fit entirely (e.g., SFT at batch size 1 uses ~30GB on OLMo-7B, while FPO at batch size 8 uses ~40GB).
Benchmark Preservation:
- FPO: Maintains off-domain benchmarks (MMLU, ARC-Challenge, HellaSwag, Winogrande) within 0.22 points of the baseline. On OLMo-2-7B, it even showed a +0.91 improvement.
- Full SFT & LoRA: Consistently degraded benchmarks. Full SFT on OLMo showed high variance, with some seeds causing severe damage (up to -8.86 points), while LoRA showed consistent degradation (-0.68 to -3.96 points).
In-Domain Performance:
- FPO achieves in-domain perplexity improvements, though smaller than full SFT or LoRA (e.g., -6.3% vs -11.4% for OLMo).
- Crucially, FPO maintains held-out text perplexity close to the baseline, whereas LoRA often degrades it significantly (e.g., +19.4% on Qwen3).
Mechanism Control (SFT-partial):
- Running standard SFT restricted only to the FPO target layers (freezing the rest) reproduced the benchmark preservation of FPO. This confirms that layer selection (late-layer adaptation) is the primary driver of benchmark safety, while FPO is the practical mechanism to achieve this regime efficiently.

5. Significance and Claims

The paper claims that FPO operationalizes the observation that late-layer gradients are largely redundant with output-layer signals. Its significance lies in:

Feasibility: Enabling domain adaptation on consumer GPUs and edge devices where full backpropagation is impossible due to memory constraints.
Stability: Providing a "benchmark-safe" regime where models adapt to new domains without forgetting general capabilities, a property not reliably achieved by full-network fine-tuning.
Efficiency: Offering a 3× throughput advantage over standard fine-tuning, making iterative development and large-scale personalization significantly faster.

The authors note limitations: FPO is currently validated on causal language models (3B–8B) with specific normalization layouts. It targets next-token cross-entropy and requires re-derivation for instruction tuning or RL objectives. The method relies on the assumption that $W_3$ changes slowly enough for the cached initial weights to serve as a valid feedback matrix for $W_1$ and $W_2$ updates.