Source-linked AI summary
The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction
Pratyusha Sharma, Jordan T. Ash, Dipendra Misra
TL;DR
Large Transformer models require substantial resources, motivating methods that reduce parameters without sacrificing capability. This paper introduces LASER, which selectively removes higher-order SVD components from specific weight matrices after training, and finds performance gains across language tasks and an embodied reinforcement-learning setting. The analysis links improvements to weakly learned, infrequent information and conflicting responses in removed components.
Problem
Transformer models are increasingly large and costly, while prior pruning work mainly sought to preserve performance after removing parameters.
Method
LASER selectively replaces specific Transformer weight matrices with low-rank approximations that remove higher-order SVD components, without additional parameters or data.
Results
LASER improves performance across five datasets and three language models, with matrices sometimes reduced to 99% below their original rank, and also benefits a decision Transformer in an embodied domain.
Takeaways & Limitations
Improvements primarily target less frequent training information, while higher-order components can encode noisy or conflicting responses whose removal helps recover correct answers.
Takeaways & Limitations
The mechanisms behind noisy higher-order components, architectural dependence, and the preference for later MLP layers remain unresolved.
Abstract
from arXiv · showhide
Transformer-based Large Language Models (LLMs) have become a fixture in modern machine learning. Correspondingly, significant resources are allocated towards research that aims to further advance this technology, typically resulting in models of increasing size that are trained on increasing amounts of data. This work, however, demonstrates the surprising result that it is often possible to significantly improve the performance of LLMs by selectively removing higher-order components of their weight matrices. This simple intervention, which we call LAyer-SElective Rank reduction (LASER), can be done on a model after training has completed, and requires no additional parameters or data. We show extensive experiments demonstrating the generality of this finding across language models and datasets, and provide in-depth analyses offering insights into both when LASER is effective and the mechanism by which it operates.
1 Introduction
Transformer models are powerful but heavily over-parameterized, motivating post-training interventions that preserve learned behavior while reducing components. LASER selectively removes higher-order weight components, often improving reasoning performance, especially in later MLP layers.
- Method: LASER replaces selected Transformer weight matrices with low-rank approximations, removing higher-order SVD components after training.The intervention targets specific matrices and layers rather than uniformly reducing the entire model.
- Findings: Performance degradation often appears only after removing well over 90% of components, while targeted reductions can instead improve accuracy on reasoning benchmarks.The reported gains extend beyond language to reinforcement learning.
- Findings: LASER improvements predominantly affect information that appears less frequently in training data, suggesting that rank reduction helps expose weakly learned facts.The paper also reports increased robustness to paraphrases on questions the original model already answered correctly.
- Mechanism: The proposed mechanism is that noisy higher-order components conflict with lower-order components, producing averaged responses that are often incorrect.Removing those components can flip answers from generic high-frequency words to the correct semantic response.
2 Related work
Prior work established that neural networks can be heavily pruned or approximated while preserving behavior, but generally expected greater intervention severity to reduce performance. This paper distinguishes LASER through targeted rank reductions that can improve predictive accuracy.
- How facts are stored: Research on stored facts indicates that information is distributed across layers, but evidence remains conflicting about how models organize and use it to construct answers.Some theories locate entity information in MLP key-value memories and describe later attention-based copying.
- Model compression: Earlier pruning and low-rank methods showed that models can lose over 90% of parameters with little accuracy drop, mainly reducing storage or inference costs.Related work also studied structured pruning and smaller networks trained to mimic larger ones.
- Low-rank approximations: Low-rank approximations typically retain the largest singular components, preserve model behavior at moderate severity, and eventually reduce performance as more components are removed.Prior reductions were commonly applied unilaterally across weight matrices.
- Novelty: This paper differs by showing that a targeted rank reduction affecting even one weight matrix can improve Transformer predictive accuracy.The related-work discussion identifies this as distinct from prior unilateral reductions.
3 Preliminaries
The preliminaries define Transformer components, matrix notation, and SVD-based rank approximation. They establish that LASER keeps large-singular-value components while removing smaller-singular-value components, with effects varying across layer types.
- Transformer Architecture: A Transformer stacks blocks that use self-attention to mix information across positions and feed-forward MLPs to process each position.The relevant matrices include attention projections and the MLP input and output matrices.
- Rank-r Approximation and SVD: SVD expresses a weight matrix through orthonormal singular vectors and descending singular values, enabling an optimal rank-r approximation by retaining the leading components.The Eckart–Young–Mirsky theorem supplies the optimal solution for a specified rank.
- Terminology: LASER calls small-singular-value components higher-order and removes them, while lower-order components correspond to large singular values and are retained.The paper’s notation uses σ_i(W) for the ith largest singular value.
- Empirical context: Rank reduction affects layer types differently: attention matrices can be reduced with little performance increase, whereas later MLP reductions can improve performance.For GPT-J on CounterFact, the transition from harm to improvement occurs around layer 20 for MLPs.
4 LAyer SElective Rank Reduction (LASER)
LASER specifies which Transformer matrix and layer to modify and how severely to reduce its rank. These interventions can be searched and composed to select modifications that yield the greatest performance benefit.
- Intervention definition: A LASER intervention is defined by parameter type, layer number, and rank reduction, determining which matrix is replaced and the severity of approximation.The parameter types include attention and MLP matrices.
- Intervention definition: The method can target an MLP input matrix at a selected layer, with the rank-k approximation controlling how much of the matrix is retained.Figure 1 illustrates an example targeting the first MLP layer within the final Transformer block.
- Search and composition: LASER throttles selected information pathways, and multiple interventions can be composed in any order.The approach searches over intervention sets and applies the modification offering the greatest benefit.
5 Experiments
Experiments show that selectively removing higher-order components can improve factual accuracy, generalize across models and tasks, and reveal patterns in which facts benefit.
- GPT-J and CounterFact: LASER improves later MLP layers most strongly, especially MLP input matrices, while attention reductions usually yield smaller gains.Subsequent experiments therefore search primarily over MLP layers.
- GPT-J and CounterFact: 24.1% versus 13.3% top-1 accuracy on CounterFact after single-layer rank reduction, without additional training or fine-tuning.The corrected set grows systematically as reduction increases.
- Which facts benefit: Facts recovered by LASER are most likely to occur infrequently in the training data.The analysis compares originally correct samples with questions corrected only after intervention.
- What higher-order components store: Higher-order components can produce generic or incorrect same-type answers, whereas retaining lower-order components can flip predictions to the correct entity.Average cosine similarity to the true answer worsens when higher-order components are used alone, motivating a conflict-resolution hypothesis.
- Generality: Across seven language-understanding datasets and three language models, severe reductions generally preserve accuracy and sometimes improve it, though the required reduction varies.Some cases improve accuracy while slightly worsening loss.
- Non-text domains: LASER solves 3% more Sokoban tasks for a decision Transformer, although non-text improvements are much smaller.The authors note that this may reflect text specificity or a need for larger Transformers.
6 Conclusion and Discussion
LASER improves question answering across models and datasets, with gains also observed in reinforcement learning, but its underlying mechanisms and scope remain incompletely understood.
- LASER improves question-answering performance across five datasets and three language models, with reductions reaching 99% of original matrix rank.The paper also reports gains for a decision Transformer in an embodied domain.
- The largest accuracy improvements correspond to information that is less common in training data, and LASER improves robustness to question paraphrases.
- Higher-order components can encode frequent words or alternate answers that overpower stable lower-order components, making LASER act as a denoising technique.
- Understanding why higher-order components accumulate noisy answers, how architecture affects LASER, and why later MLP layers matter requires further study.
A Dataset Details
The paper evaluates LASER across factual, linguistic, reasoning, truthfulness, and embodied-planning settings using several datasets and a decision Transformer Sokoban environment.
- The evaluation includes CounterFact, QA Wiki Data, PILE, HotpotQA, FEVER, Bias in Bios, BigBench-Epistemic Reasoning, and TruthfulQA.
- CounterFact contains knowledge tuples derived from Wikidata, handwritten prompt templates, and suggested edits to represented facts.
- PILE is an approximately 1TB language-modeling dataset used to study LASER on GPT-J’s training distribution and entity frequency.
- HotpotQA uses 7,100 validation examples to identify LASER parameters and 90,000 training examples for evaluation.
- Sokoban Details: Sokoban tests long-horizon planning by requiring a five-layer decision Transformer trained on 106 optimal episodes to move boxes to targets without trapping them.The maximum return is set to 10.
C.2 How much reduction is too much?
Rank reduction typically improves performance up to a task- and layer-dependent point, after which performance worsens; completely removing a layer is intermediate between full-rank and LASER treatment.
- Performance first improves monotonically with increasing rank reduction, then worsens, and this pattern is observed consistently across interventions.The improvement threshold varies with layer type and layer number.
- Completely removing a layer can outperform retaining its full-rank matrix but performs worse than applying a low-rank approximation.
- Figure 7 compares effective ranks of GPT-J MLP matrices where LASER helps, showing that reductions extend beyond the matrices’ computed effective ranks.
C.3 Does LASER select the same layer for different tasks?
LASER’s best-performing layer varies across tasks, while task gains can coexist with a slight worsening of language-modeling perplexity.
- Different tasks achieve their maximum improvements from LASER applied to different layers of the model.
- GPT-J models selected as best across tasks reduce matrices in different layers, rather than sharing one universally optimal layer.
- LASER slightly worsens GPT-J perplexity on PILE even when it improves the task-specific performance being evaluated.The paper leaves the meaning of this perplexity change for future study.
C.5 Final LASER search results
The final LASER search finds that optimal improvements typically come from reducing MLP input matrices in later Transformer layers, sometimes to just 1% of their original rank.
- Optimal LASER improvements typically come from reducing the MLP Input matrix in later Transformer layers.The search selects parameters maximizing validation accuracy for each model–dataset pair.
- The rank is sometimes reduced to 1% of the original matrix rank.
D Alternate Pruning Methods
The paper compares LASER with alternative pruning and examines how reduction severity and layer location affect performance. Benefits depend on the dataset–model pair, while excessive reduction can eventually hurt accuracy.
- GPT-J on CounterFact improves similarly under layer-selective absolute weight pruning and LASER.Absolute pruning removes lower-magnitude weights, and further study of its connection to LASER is left for future work.
- Performance improves with large reductions but begins worsening after a point, including a dip at 99.95% reduction.This result is reported for GPT-J on CounterFact using top-10 accuracy.
- GPT-J on CounterFact reduces an MLP matrix from rank 4096 to rank 4, approximately 99% of its original rank.
- The largest LASER benefit comes from different layer numbers across GPT-J dataset–model pairs, although later MLP layers often help most.
- Pruning later MLP layers also increases GPT-J’s accuracy on CounterFact.
E.1 Dataset Processing Details
The evaluation processes each dataset separately, reserves 20% of processed data for LASER hyperparameter selection, and uses task-specific prompting and accuracy procedures across datasets and model types.
- Each dataset is processed separately, with 20% of the processed data used to select LASER hyperparameters by validation accuracy.The remaining data evaluates the selected model configuration.
- The processed benchmarks include CounterFact, Hotpot, Fever, Bios Gender, Bios Profession, BigBench Epistemic Reasoning, TruthfulQA, and BigBench Wikidata QA.Dataset-specific filtering, prompting, and sample counts vary across these tasks.
- TruthfulQA is converted into separate true-or-false question–answer samples, producing 5,882 processed examples.
- Open-ended tasks generally use greedy generation, while tasks with small label sets use the highest-probability label for classification accuracy.The evaluation lowercases and strips whitespace before comparing generated answers, with task-specific generation lengths.
- Roberta generates through masked tokens, using a fixed number for responses but matching the answer length when computing log loss.
- Generation lengths and label scoring are adapted to each task, including up to 15 tokens for HotPotQA and up to 10 for BigBench Wikidata QA.