The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction — part 2/4

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

Given a matrix W ∈ ℝ m × n W\in\mathbb{R}^{m\times n} and r ∈ ℕ r\in\mathbb{N} , a rank- r r approximation problem requires finding a matrix W ^ \hat{W} that minimizes ‖ W − W ^ ‖ 2 \|W-\widehat{W}\|_{2} and satisfies 𝚛𝚊𝚗𝚔 ⁡ ( W ^ ) ≤ r {\tt rank}\left(\widehat{W}\right)\leq r . Eckart–Young–Mirsky theorem provides an optimal solution of this problem using Singular Value Decomposition (SVD) ( Eckart and Young, 1936 ) . Formally, an SVD of a matrix W W is given by W = U ​ Σ ​ V ⊤ W=U\Sigma V^{\top} where U = [ u 1 , u 2 , ⋯ , u m ] ∈ ℝ m × m U=[u_{1},u_{2},\cdots,u_{m}]\in\mathbb{R}^{m\times m} and V = [ v 1 , v 2 , ⋯ , v n ] ∈ ℝ n × n V=[v_{1},v_{2},\cdots,v_{n}]\in\mathbb{R}^{n\times n} and Σ ∈ ℝ m × n \Sigma\in\mathbb{R}^{m\times n} . The column vectors of U U and V V constitute an orthonormal basis of ℝ m \mathbb{R}^{m} and ℝ n \mathbb{R}^{n} respectively, and Σ \Sigma is a diagonal matrix whose diagonal entries are given by the singular values of W W in descending order. One can also express the SVD of W W as W = ∑ i = 1 min ⁡ { m , n } σ i ↓ ​ ( W ) ​ u i ​ v i ⊤ W=\sum_{i=1}^{\min\{m,n\}}\sigma^{\downarrow}_{i}(W)u_{i}v^{\top}_{i} . According to Eckart–Young–Mirsky theorem, the matrix W ^ = ∑ i = 1 r σ i ↓ ​ ( W ) ​ u i ​ v i ⊤ \widehat{W}=\sum_{i=1}^{r}\sigma^{\downarrow}_{i}(W)u_{i}v^{\top}_{i} is an optimal solution to the rank- r r approximation problem for any given desired rank r ≤ min ⁡ { m , n } r\leq\min\{m,n\} .

In this work, we will use the word higher-order components to refer to entries in the SVD corresponding to the components with smaller singular values. These components are removed by 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} . The term lower-order components is used to refer to singular vectors corresponding to large singular values. These components are kept in a low-rank approximation of the matrix.

## 4 LA yer SE lective R ank Reduction ( 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} )

In this section, we formally describe the 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} intervention. A single-step 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} intervention is defined by three quantities ( τ , ℓ , ρ \tau,\ell,\rho ), consisting of a parameter type τ \tau , layer number ℓ \ell , and rank reduction ρ \rho . These values together describe which matrix will be replaced by their low-rank approximation and how severe the approximation will be. The parameter type categorizes the matrix type in which we will intervene. We focus on the matrices in 𝒲 = { 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\} which consist of the matrices in the MLP and attention layers. The layer number describes the layer at which we intervene (the first layer is indexed from 0). E.g., the Llama-2 has 32 layers and so ℓ ∈ { 0 , 1 , 2 , ⋯ 31 } \ell\in\{0,1,2,\cdots 31\} . Finally, ρ ∈ [ 0 , 1 ) \rho\in[0,1) describes what fraction of the maximum rank should be preserved upon doing its low-rank approximation. For example, let τ = U i ​ n ∈ ℝ d × d \tau=U_{in}\in\mathbb{R}^{d\times d} , then the maximum rank of this matrix is d d . We replace it with a rank ⌊ ρ ⋅ d ⌋ \lfloor\rho\cdot d\rfloor -approximation.

Figure 1 shows an example of 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} . In this figure, we have τ = U i ​ n \tau=U_{in} and ℓ = L \ell=L indicating that we update the weight matrix in the first layer of MLP in the Transformer block of the L t ​ h L^{th} layer. The other parameter (not shown in the figure) controls the k k in the rank- k k approximation.

𝙻𝙰𝚂𝙴𝚁 {\tt LASER} throttles the flow of certain information in the network, which surprisingly can produce significant performance benefits. These interventions can also be easily composed—we can apply a set of interventions { ( τ i , ℓ i , ρ i ) } i = 1 m \{(\tau_{i},\ell_{i},\rho_{i})\}_{i=1}^{m} in any order. The 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} approach is to simply search over interventions of this type, and to exercise the modification that offers the greatest benefit. There are many other ways in which one can combine these interventions, however, and we defer a this to future work.

## 5 Experiments

This section studies the consequences of 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} throughout various layers of the Transformer architecture. We first perform a motivating analysis of the CounterFact ( Meng et al., 2022 ) question-answering dataset in conjunction with a pretrained GPT-J model ( Wang and Komatsuzaki, 2021 ) , and investigate the performance of the model and its variability as we search over potential interventions. Following that, we look at the effect of 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} across different models, datasets and modalities.

##### GPT-J, CounterFact and PILE.

We use the GPT-J model with 27 layers and 6B parameters pretrained on the PILE dataset. The first part of our analysis focuses on GPT-J, largely because its training data is publicly available. We evaluate the model’s behavior on the CounterFact dataset, which consists of samples organized as (subject, relation, answer) tuples and three paraphrased prompts for each question. For example, (Danielle Darrieux, mother tongue, French).

### 5.1 A Thorough Analysis with GPT-J on the CounterFact Dataset

Figure 2 shows the result of applying various amounts of rank reduction to each matrix in the Transformer architecture on the classification loss for this dataset. These plots are grouped, such that each sub-figure corresponds only to the indicated type of weight matrices. Note that each Transformer layer consists of a small, two-layer MLP. The constituent input and output matrices are shown separately. Different colors indicate different percentages of removed components.

The attention plots in this figure exemplify what is already known about these models: weight matrices can be drastically reduced without much degradation in model performance. The more interesting result, however, is in the MLP layers. Here, not only can matrices be rank-reduced without degrading classification performance, but large performance improvements are possible by reducing later layers of the model. This trend is most stark in the input matrix of the MLP. While there are gains with 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} in the attention layers too, the benefits are typically smaller. In the section that follows, we demonstrate the effectiveness of 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} across a wide array of datasets and Transformer models. Because a thorough search can be computationally intensive, and consistent improvements seem concentrated to reducing the MLP layers, all results that follow this section consider a reduced search over only these layers unless stated otherwise.

##### Improved accuracy and robustness to paraphrases.

The CounterFact dataset is used to test the model’s factual knowledge of data from Wikipedia. Since GPT-J is trained on PILE, whose contents include Wikidata, different facts in CounterFact are part of the model’s training data, albeit in different quantities. As all answers are a single token in this setting, we compute top-k accuracy based on whether the correct answer is in the top-k predicted tokens. As seen in Figure 2 and Table 1 , we find that the model’s top-1 accuracy on facts in CounterFact increases from 13.3% to 24.1% when reductions are done on a single layer. It is important to note that these improvements are a result of rank-reduction alone, and do not involve any further training or fine-tuning of the pre-trained GPT-J model. Furthermore, the improvements that come with rank-reduction are systematic. The set of datapoints that the model gets correct only grows with increasing amounts of reduction as opposed to a random movement of datapoints into and out of the set or correct items; if a model gets an answer right with a certain amount of rank reduction ( x x ), the model continues to get the answer correct for larger rank reductions ( y y where y > x y>x ). We evaluate the model’s robustness to paraphrases by computing the percentage of datapoints where the model gets all paraphrases of a given question correct. For datapoints that the model already gets correct, the model’s robustness to paraphrases also improves with 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} by roughly 24.8 percentage points.

##### Effect on language modeling and fluency.

While the model’s factuality improves, does the reduction affect the model’s performance on other metrics? To understand this, we evaluate the model’s perplexity, i.e., its original training objective, on its training data. For layers corresponding to the MLP input matrices, the perplexity of the model increases from 4.8 to 5.0, showing that the language modeling objective is indeed slightly effected. For the MLP output layers, the perplexity of GPT-J on PILE increases from 4.8 to 4.9 with 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} . It may be possible to fix this small degradation by calibrating the temperature of the model.

##### Composing reductions across layers.

We find that even further improvements in the model’s performance can be made by performing different amounts of rank reduction across several layers. This is done by greedily searching over ( τ , ℓ , ρ \tau,\ell,\rho ) starting from the largest ℓ \ell and smallest ρ \rho . To speed things up, here we do this search only over MLP layers, as this is where the largest improvements are typically found. Consistent with other experiments, the search is done on a validation set, and results are reported on the test set. On CounterFact, the 0-1 accuracy of the base GPT-J model is 13.1%. After doing the best single-step 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} the model’s accuracy improved to 24.0%. Performing 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} across different layers improved the top-10 accuracy to 29.2%, a 5.2% absolute improvement in accuracy over performing 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} on a single layer. The results of the combinatorial search across different ℓ \ell and ρ \rho values can be seen in Figure 4 .

#### 5.1.1 Which facts in the dataset are recovered by rank reduction?

To understand this phenomenon, we look at the questions correctly answered after 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} and the effect of how often the information associated with the question appears in the training data. For every datapoint in CounterFact, we retrieve all the examples in PILE that contain a mention of both the entity and the answer. We then compute how often information associated with each evaluation question appears in the training data. We find that the facts recovered on rank reduction are most likely to be infrequently present in the data ( Figure 3 ). Here, “Originally correct” describes samples that are correctly classified even without any intervention. “Answer-corrected” refers to questions the model gets correct only after intervening with 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} .

#### 5.1.2 What are higher-order components storing?

We saw above how retaining the lower-order components improves model performance on the task of open-ended question answering. We find that for the task of question answering, improvements often come on questions whose answers are supported by less frequently occurring data in the training set. While it is clear that eliminating the higher-order components “denoises” the model and helps recover “hidden,” less-frequent information, it is not clear what the higher-order components are representing such that that their removal improves performance. This section studies this question using the CounterFact dataset and GPT-J.

To understand what higher-order components are representing, we approximate the final weight matrix using its higher-order components (as opposed to approximating it using its lower-order components as done by 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} ) as shown in Figure 5 (a). Following this, we analyze how the model’s behavior changes on data points that GPT-J originally gets incorrect but are flipped to being correct upon performing 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} .

First, we note that when the original, unmodified model does not answer these questions correctly, it often responds with common words, such as “a,” “the,” “of,” and other highly frequent tokens. After performing 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} , where we retain only the top-k components, the model’s answers to these questions flip from generic words to the correct entity. For the same datapoints, when we approximate the model by instead retaining the higher-order components, we find that the model either predicts incorrect entities that are of the same semantic type as the correct answer or high-frequency tokens such as “a,” “the,” and “of,” as shown in Figure 5 (c). However, as we systematically include the lower-order components, the model’s output changes to predicting frequent tokens. To investigate this systematic degradation, we measure the average cosine similarity of the “true” answer with respect to the predicted answer when the matrix is approximated with different amounts of higher-order components, as shown in Figure 5 (b). The average cosine similarity between the predicted answer worsens, demonstrating this effect.

We hypothesize that these matrices often encode multiple conflicting responses, and that when all components are used they clash to produce a generic token. Removing the higher-order components, which anecdotally appear to often capture incorrect responses of the correct type, resolves this internal conflict and allows the model to respond accurately.

### 5.2 How generally does this hold?

We evaluate the generality of our findings on 3 different LLMs for several language understanding tasks.

##### Natural Language Understanding Tasks.

We evaluate model performance before and after 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} on seven datasets, including CounterFact ( Meng et al., 2022 ) , HotPotQA ( Yang et al., 2018 ) , FEVER ( Thorne et al., 2018 ) , Bias in Bios ( De-Arteaga et al., 2019 ) [Gender and Profession], TruthfulQA ( Lin et al., 2021 ) , BigBench-Epistemic Reasoning ( Bowman et al., 2015 ) and BigBench-WikidataQA. These datasets evaluate different aspects of language understanding problems. CounterFact, Fever, and Bigbench-Wiki data test a model’s world knowledge and factuality. Bias in Bios benchmarks model bias by predicting the gender and profession of a person given a short biography. We define Bios Gender as the gender prediction problem in Bias in Bios, and Bios Profession as the profession prediction problem. HotPotQA provides a more challenging open-ended question answering task with long answers containing many tokens. The Epistemic Reasoning dataset from Big Bench Hard (BBH) tests a model’s logic and reading comprehension. Finally, TruthfulQA tests an LLM’s truthfulness.We use 20% of the dataset as validation set and select the best 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} hyperparameters ( τ , ℓ , ρ ) (\tau,\ell,\rho) using this validation set. We report results on the remaining 80% of the dataset with the chosen hyperparameter. The models used for the task of question answering include, Roberta ( Liu et al., 2020 ) , GPT-J (6B) ( Wang and Komatsuzaki, 2021 ) , and LLAMA2 (7B) ( Touvron et al., 2023 ) . Details regarding datasets and how they were used can be found in Appendix A .

##### Evaluation metrics.

For each of these tasks, we evaluate the model’s performance using (i) generation accuracy. We generate a sequence of N N tokens using the LLM and then report 1 if the answer text is in the generated text and 0 otherwise, (ii) classification accuracy . If the answer lies in a small set of potential values, like in a standard classification problem, we consider a response correct if it puts more probability mass on the correct answer than on any of the other candidates, and (iii) loss. We report the log-loss on held-out data. For datasets with a small set of possible labels, we report the accuracy ( acc ) using classification accuracy, whereas for others we use the generation accuracy.