The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction β€” part 1/4

Cleaned-text code points [0, 15285). Bibliography omitted here; source.md includes it.

# The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction

###### Abstract

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 1 1 1 Higher-order components are singular vectors with smaller singular values. of their weight matrices. This simple intervention, which we call LAyer-SElective Rank reduction ( π™»π™°πš‚π™΄πš {\tt 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 π™»π™°πš‚π™΄πš {\tt LASER} is effective and the mechanism by which it operates 2 2 2 Code and website: https://pratyushasharma.github.io/laser/ .

## 1 Introduction

Since their original release, Transformer-based LLMs have been shown to be remarkably proficient on a wide array of important machine learning tasks. Their underlying Transformer architecture has become state-of-the-art for modeling and reasoning about natural language, and has shown promise in domains such as computer vision ( Dosovitskiy et al., 2020 ) and reinforcement learning ( Chen et al., 2021 ) as well.

Contemporary instantiations of Transformer architectures are infamously large, typically requiring tremendous compute resources for both training and inference. This is by design, as Transformers trained with more parameters or data are demonstrably more capable than their slimmer predecessorsβ€”often by a significant margin ( Brown et al., 2020 ; Touvron et al., 2023 ) . Still, a growing body of work suggests that Transformer-based models, and neural networks more generally, do not require all fitted parameters to retain their learned hypotheses. While it seems helpful to be massively over-parameterized at train time ( Hinton et al., 2015 ; Bengio et al., 2005 ) , it is well-known that these models can be drastically pruned before inference; neural networks can often have well over 90% of their weights removed without any significant degradation in performance ( Frankle and Carbin, 2018 ) . The discovery of this phenomenon bolstered interest around the relationship between generalization and over-parametrization ( Zhang et al., 2017 ) , and spawned research in developing pruning strategies that lend themselves to efficient model inference ( Molchanov et al., 2016 ) .

This paper presents a surprising finding, that careful pruning done at specific layers of Transformer models can produce significant boosts in performance on some tasks. We describe LAyer SElective Rank reduction ( π™»π™°πš‚π™΄πš {\tt LASER} ), an intervention that removes higher-order components of learned weight matrices as identified by singular value decomposition. This reduction is performed in specific weight matrices and layers of the Transformer model. In line with previous work, we find that many such matrices can be significantly reduced, and that performance degradation is often not observed until well over 90% of components are entirely removed. However, unlike what is found in previous work, we find that these reductions can produce drastic improvements in accuracy, as measured by various well-studied reasoning benchmarks in NLP. Even better, this discovery appears to not be limited to natural language, with performance gains also found in reinforcement learning.

This paper analyzes the relationship between the model’s training data and samples that benefit from π™»π™°πš‚π™΄πš {\tt LASER} . We find that improvements in the model’s performance predominantly come on information less frequently present in the model’s training dataset, suggesting that π™»π™°πš‚π™΄πš {\tt LASER} offers a kind of denoising procedure that makes weakly learned facts accessible. We separately observe that π™»π™°πš‚π™΄πš {\tt LASER} affords increased robustness to paraphrases on previously correct questions.

Additionally, we attempt to reason about what is being stored in high-order components, such that their removal boosts performance. For questions correctly answered only after π™»π™°πš‚π™΄πš {\tt LASER} , in the absence of interventions, the original model predominantly responds with high-frequency words such as β€œthe,” β€œof,” etcβ€”generations that are not even the same semantic type as the correct answer. However, after some amount of rank reduction, the model’s answer flips to be correct. To understand this, we look at what the remaining components encode on their own; we approximate the weight matrix using only its higher-order singular vectors. We find that these components describe either a different response of the same semantic category as the correct answer or generic high-frequency words. Seemingly, when the noisy, higher-order components are combined with the low-order components, their conflicting responses produce a sort of β€œaverage answer,” which is likely to be incorrect.

Figure 1 visualizes the Transformer architecture and the procedure followed by π™»π™°πš‚π™΄πš {\tt LASER} . Here, the weight matrix of a Multi-Layer Perceptron (MLP) at a specific layer is replaced with its low-rank approximation.

## 2 Related work

To our knowledge, this paper is the first to identify that carefully selected rank reductions can boost Transformer performance. Still, there is a wide array of works that study related questions, including how facts are stored in LLMs and how to best compress neural networks.

##### How facts are stored.

Studies probing model representation for the presence of select properties of entities ( Ettinger et al., 2016 ; Adi et al., 2016 ; Hupkes et al., 2018 ; Conneau et al., 2018 ) show that models store factual information across different layers and Lee et al. (2023) show that model robustness to distribution shift can be improved by fine-tuning select layers. However, there is conflicting evidence on how this information is organized and utilized in constructing answers in large language models. Some theories posit that information about different entities is locally stored as two-layer, key-value memory in MLP sections of Transformer models ( Geva et al., 2021 ) , which are thereafter copied over through latter layers by the self-attention modules ( Elhage, 2021 ) . Meng et al. (2022) propose a procedure to trace and edit local, entity-specific information to map to distinct β€œimpossible” outputs, supporting the locality theory. These theories are further supported by the phenomenon of β€œearly exiting,” where the representation at an intermediate layer can be directly used with the terminal head of the model to correctly generate an output ( Zhao et al., 2021 ) . In contrast, Hase et al. (2023) has observed that information about some of the same entities or entity relations can be modified by editing a variety of layers in the model architecture, and therefore, that facts are stored across layers in a fragmented fashion. This paper makes no specific claims regarding locality, but rather demonstrates that higher-order components of a weight matrix can introduce noise in decision making, and considering only lower-order components may make correct answers accessible.

##### Model compression.

Neural network pruning methods have found that models could be significantly pruned (often removing over 90% of parameters) with very little drop in accuracy, significantly reducing the storage requirements of the model ( LeCun et al., 1989 ; Hassibi and Stork, 1992 ; Han et al., 2015 ; Li et al., 2016 ; Frankle and Carbin, 2018 ) . There have also been approaches that prune these models in a structured manner, to facilitate improvements in inference time ( Molchanov et al., 2016 ) . The existence of sparse sub-networks ( Frankle and Carbin, 2018 ; Hoefler et al., 2021 ) has been found to be true for convolutional, fully connected, and Transformer models ( Lv et al., 2023 ; Murty et al., 2022 ) . While Jin et al. (2022) find that model generalization can be improved by pruning and then refitting parameters, generalization improvements are only observed upon model retraining. To our knowledge, model pruning techniques have always done a unilateral reduction across all parameters, without targeting any specific layersβ€”leading to predictive performance either staying fixed or decreasing ( Frankle and Carbin, 2018 ) . In this work, however, we find that the effect of reduction in accuracy is non-uniform across different layer types, and a model’s generalization can be improved by selective pruning alone; no additional training is necessary. Roughly, we find that performance degradation can be produced by rank-reducing early layers, while significant performance benefits are typically available by pruning later layers.

##### Low-rank approximations of weight matrices.

Most pruning methods reduce parameters in order of their absolute magnitude ( Frankle and Carbin, 2018 ) . An alternative approach, however, is to reduce the rank of its constituent weight matrices, keeping the top k k components found by SVD. While matrices of neural models, including Transformer models, have been found to be well-approximated using this approach, where markedly reduced versions of the model can preserve its behavior, research has shown that performance eventually declines as the severity of the intervention increases ( Lv et al., 2023 ; Hajimolahoseini et al., 2021 ; Yu et al., 2017 ) . Note that these reductions are typically done unilaterally, removing the same number of components in every weight matrix in the model. In contrast to these findings, we show that a targeted rank reduction, even affecting just a single weight matrix, can offer benefits to the predictive accuracy of Transformers.

##### Model distillation and low-rank training.

Ba and Caruana (2014) and Hinton et al. (2015) have trained smaller networks to mimic the behavior of larger networks, suggesting neural networks might be significantly over-parametrized and can be replaced with leaner alternatives. To our knowledge, no report of an improvement in the model’s predictions as a consequence of this procedure has been shown. ( Yang et al., 2020 ) have enforced low-rank-ness of weight matrices for the purposes of memory efficiency, but the resulting models fail to achieve performance equivalent to their overparametrized counterparts. The result suggests that overparametrization is helpful for the identification of well-generalizing parameters by SGD ( Bengio et al., 2005 ; Hinton et al., 2015 ; Zhang et al., 2017 ) .

## 3 Preliminaries

Here we review basic notation and describe the core components of our study.

##### Maths Notation.

We use ℝ \mathbb{R} to denote real numbers, β„• \mathbb{N} to denote natural numbers, small letters such as v ∈ ℝ d v\in\mathbb{R}^{d} to denote a d d -dimensional vector, and capital letters such as W ∈ ℝ m Γ— n W\in\mathbb{R}^{m\times n} to denote a matrix of size m Γ— n m\times n . We use β€– v β€– 2 \|v\|_{2} to denote the Euclidean norm of a vector v v and β€– W β€– 2 \|W\|_{2} to denote the spectral norm of a matrix W W . We use [ N ] [N] to denote the set { 1 , 2 , β‹― , N } \{1,2,\cdots,N\} . We will use πš›πšŠπš—πš” ⁑ ( W ) {\tt rank}(W) to denote the rank of a matrix W W and Οƒ i ↓ ​ ( W ) \sigma^{\downarrow}_{i}(W) to denote its i t ​ h i^{th} largest singular value.

##### Transformer Architecture.

We provide a concise description of vanilla Transformer architecture that is relevant to our analysis. A Transformer architecture can be thought of as L L layers of Transformer blocks. The l t ​ h l^{th} block maps a sequence of T T -length vector sequence ( h 1 ( l βˆ’ 1 ) , β‹― , h T ( l βˆ’ 1 ) ) (h^{(l-1)}_{1},\cdots,h^{(l-1)}_{T}) to another T T -length vector sequence ( h 1 ( l ) , β‹― , h T ( l ) ) (h^{(l)}_{1},\cdots,h^{(l)}_{T}) , where all vectors are d d -dimensional. This transformation is accomplished using two sequential steps: a self-attention mechanism to mix information across time steps, and a feed-forward network to process information within each time step. We describe a basic version of these transformations for a fixed l t ​ h l^{th} layer and drop the superscript ( l βˆ’ 1 ) (l-1) for clarity. 3 3 3 Various Transformer models often have small differences in how these transformations are implemented. Our goal is not to provide a full survey of these details but to capture essential terminology for our results.

A single-head self-attention mechanism first maps each vector h i h_{i} to a query vector q i = W q ​ h i q_{i}=W_{q}h_{i} , a key vector k i = W k ​ h i k_{i}=W_{k}h_{i} and a value vector v i = W v ​ h i v_{i}=W_{v}h_{i} where W q , W k , W v ∈ ℝ d Γ— d W_{q},W_{k},W_{v}\in\mathbb{R}^{d\times d} are layer-specific weight matrices. We then compute attention probabilities p ⁑ ( j ∣ i ) = exp ⁑ ( q i ⊀ ​ k j / d ) βˆ‘ l = 1 T exp ⁑ ( q i ⊀ ​ k l / d ) p(j\mid i)=\frac{\exp(q_{i}^{\top}k_{j}/\sqrt{d})}{\sum_{l=1}^{T}\exp(q_{i}^{\top}k_{l}/\sqrt{d})} for every i , j ∈ [ T ] i,j\in[T] . These are used to compute the attention vector z i = βˆ‘ j = 1 T p ⁑ ( j ∣ i ) ​ v j z_{i}=\sum_{j=1}^{T}p(j\mid i)v_{j} . A k k -head self-attention computes a set of k k attention vectors by using different linear transformations for key, query, and value, and then concatenates these attention vectors. These k k -separate linear transformations for key, query, and value can all be absorbed into their respective matrices W q ∈ ℝ d Γ— d ​ k , W k ∈ ℝ d Γ— d ​ k W_{q}\in\mathbb{R}^{d\times dk},W_{k}\in\mathbb{R}^{d\times dk} and W v ∈ ℝ d Γ— d ​ k W_{v}\in\mathbb{R}^{d\times dk} . Finally, the self-attention mechanism outputs u i = z i ​ W o + h i u_{i}=z_{i}W_{o}+h_{i} using a projection matrix W o ∈ ℝ d ​ k Γ— d W_{o}\in\mathbb{R}^{dk\times d} .

The feed-forward step applies a 2-layer multi-layer perception (MLP) ψ : ℝ d β†’ ℝ d \psi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d} to each vector u i ∈ ℝ d u_{i}\in\mathbb{R}^{d} separately. The MLP typically has a ReLU or GELU activation function ( Hendrycks and Gimpel, 2016 ) and in some models such as Llama, the bias of linear layers is set to 0. We denote the weight matrices of the first and second linear layers of this MLP by U i ​ n U_{in} and U o ​ u ​ t U_{out} respectively. The output of this l t ​ h l^{th} Transformer block is then given by h i ( l ) = ψ ⁑ ( u i ) + u i h^{(l)}_{i}=\psi(u_{i})+u_{i} .

In summary, a Transformer architecture has the following weight matrices 𝒲 = { W q , W k , W v , W o , U i ​ n , U o ​ u ​ t } \mathcal{W}=\left\{W_{q},W_{k},W_{v},W_{o},U_{in},U_{out}\right\} for each layer, in addition to the embedding matrix for embedding input tokens, a projection weight matrix applied after the final layer before taking softmax, and all weight matrices associated with layer normalization. In our work, we will focus primarily on the matrices in 𝒲 \mathcal{W} and intervene by modifying them.

##### Rank- r r Approximation and SVD.