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

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

We test the generality of this result by evaluating a collection of language models on different benchmarks. As seen in Table 1 , we find that even severe reductions result in no deterioration in the model’s accuracy and can lead to improvements in their performance. The amount of reduction required differs from model to model.

### 5.3 Non-text domains

To understand if this phenomenon is effective outside of question answering in the textual domain, we evaluate the effect of rank reduction on a reinforcement learning agent.

##### Policy learning.

For Policy learning, we evaluate the effect of 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} on a decision Transformer model trained on the game of Sokoban and evaluated on the same game. This is a challenging planning problem where the agent has to move and push several blocks to holes. The task is completed when all blocks are on top of holes. The input to the decision Transformer is the visual state of the environment at a given state, and the output is the low-level action. We find that for a decision Transformer trained on Sokoban, models solved 3% more tasks with 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} ( Table 2 ). Details of the experiment can be found in Appendix B .

Although the improvements are much smaller, they are consistent despite the severity with which reductions are performed. This can be because the phenomenon is either text-specific or requires a large enough Transformer model.

## 6 Conclusion and Discussion

This paper describes 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} , a phenomenon where performing a low-rank approximation of specific layer types at specific layers of the transformer block can improve the performance of LLMs on the task of question answering. We find this to be true across five different datasets and three different language model models. Furthermore, the resulting 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} reductions are extreme. The matrices are reduced at times to 99% of their original rank, which is much lower than their effective rank ( C.1 ). However, despite extreme reductions, the performance of the model on tasks continues to improve. We also observe performance gains for a decision Transformer in an embodied domain. We find that the largest improvements in the model accuracy correspond to information that is less common in the training data and that 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} jointly makes the model more robust to paraphrases of the questions. We further found that the higher-order components of some of these matrices encode either high-frequency words or alternate answers of the same semantic type as the correct answer. These noisy, higher-order components can overpower the stable lower-order components and result in the model answering questions incorrectly. In these cases, performing 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} acts as a denoising technique and reduces the internal conflicts in potential responses.

Despite this analysis, the success of 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} requires further study. Learning (i) why higher-order components in weight matrices accumulate noisy answers in the course of training, (ii) the effect of model architecture and other structural choices on the occurence of this phenomenon and (iii) why this is specifically true for later layers in the MLP is important to not only for our understanding of the success of 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} , but for understanding the behavior of large language models more generally.

## Acknowledgements

This work was done when PS was an intern at Microsoft Research, New York City, with DM and JTA. The authors would like to thank Minyoung Huh, Shikhar Murty, Han Guo, Cyril Zhang, David Bau, Jacob Andreas, Antonio Torralba, and John Langford for helpful discussions and paper feedback.

## Appendix

## Appendix A Dataset Details

##### CounterFact.

The CounterFact dataset is derived from the PARAREL dataset [ Elazar et al., 2021 ] and contains knowledge tuples of the kind t c = ( s , r , o c ) t^{c}=(s,r,o^{c}) , where s s is the subject, r r is the relation and o o is the object. These tuples are constructed using entities listed in Wikidata. The datapoints are accompanied by handwritten prompt templates for each category. The CounterFact dataset also contains suggested edits to the true facts represented in the dataset. For this study, the set of counterfactual edits are not used.

##### PILE.

The PILE dataset is an approximately 1TB language modeling dataset that was used to pre-train GPT-J. It contains text from 22 smaller datasets, including Wikipedia, OpenWebText2, and StackExchange, to name a few. The PILE dataset was used to study the effect of 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} on the behavior of the model on the original training data distribution. For the study on quantifying the occurrences of entities in the training data, the training data split of PILE was used. However, the measure of change in perplexity of the model after 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} was measured on the validation split of the dataset.

##### HotpotQA.

We use the HotPotQA dataset available on HuggingFace. An example question is “ What are the names of the current members of American heavy metal band who wrote the music for Hurt Locker The Musical? " and the answer is “ Hetfield and Ulrich, longtime lead guitarist Kirk Hammett, and bassist Robert Trujillo ". We use the validation split of this dataset for identifying 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} parameters, containing 7,100 datapoints, and the train split, containing 90k datapoints, for evaluation.

##### FEVER.

The FEVER (Fact Extraction and Verification) dataset consists of claims and label characterizing each as either true or false. The goal is to predict the label given the claim. It consists of 185,445 such claims and is constructed using data from Wikipedia. An example claim is “ Roman Atwood is a content creator. ” and the label is “ true .”

##### Bias in Bios.

Bias in Bios is a dataset that consists of descriptions of people, and the task is to correctly identify the gender and occupation of the person. An example datapoint is “ Nancy Lee graduated from Lehigh University, with honours in 1998. Nancy has years of experience in weight loss surgery, patient support, education, and diabetes ,” the gender label is “ Female ” and the profession is “ Nurse .”

##### BigBench-Epistemic Reasoning:

The epistemic reasoning dataset tests a model’s ability to make entailment decisions from a pair of sentences. Specifically, epistemic reasoning assesses a given model’s understanding of “theory of mind,” which involves the capacity to analyze the cognitive states of various agents. A sample consists of a text containing a premise and hypothesis and a label which is either “entailment” or “non-entailment.” An example datapoint is “Premise: Emma knows that James thinks that there is milk in the fridge. Hypothesis: James thinks that there is milk in the fridge.” The label here is “entailment.”

##### TruthfulQA:

TruthfulQA evaluates the extent to which a model replicates human falsehoods. The benchmark contains 817 questions that span 38 categories, including health, law, finance and politics. The dataset emphasizes that solely increasing the scale of models may not be as effective in enhancing truthfulness as fine-tuning with training objectives that go beyond imitating text from the web. An example datapoint is “What happens to you if you eat watermelon seeds? [‘Nothing happens,’ ‘You eat watermelon seeds,’ ‘The watermelon seeds pass through your digestive system,’ ‘You will not digest the watermelon seeds,’ ‘The watermelon seeds will be excreted,’ ‘You grow watermelons in your stomach,’ ‘You get sick,’ ‘You have bad dreams,’ ‘You die,’ ‘You get indigestion,’ ‘You fall unconscious,’ ‘You digest the watermelon seeds’]”.

##### QA Wiki Data:

The task tests a model’s ability to answer open domain, cloze-style question answering on randomly sampled Wikipedia triplets. This evaluates world knowledge on a large collection of facts and information from a knowledge graph extracted from Wikipedia data. An example datapoint from this dataset is “Gabon shares a border with Cameroon .”

## Appendix B Details of the Decision Transformer Domain

##### Sokoban Details.

We show an image of the Sokoban task in Figure 6 . Sokoban is a warehouse-keeping transportation game that requires long-horizon reasoning and planning over multiple time steps. The task is to move all boxes to their target locations without getting trapped. We use the Gym Sokoban environment Schrader [2018] , and train a 5-layer decision transformer model using 10 6 10^{6} optimal episodes of the game. In our setting, the maximum return of the game is set to 10.

## Appendix C Extended Analysis

### C.1 Are weight matrices already low-rank?

As seen in Figure 7 , we find that 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} approximated matrices with their low-rank approximations much beyond their effective rank as computed by [ Roy and Vetterli, 2007 ] . To study this, we computed the effective rank of the MLP matrices for which 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} helps for GPT-J model using the method described by Roy and Vetterli [2007] . The plot shows that although matrices of the later layer have a lower effective rank than the earlier layers, the computed effective rank is significantly larger than the reduction % until which 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} helps.

### C.2 How much reduction is too much?

We see that for many of the matrices in cases where reduction helps, with increasing amounts of rank-reduction, the model first monotonically improves before it starts to worsen, as seen in Figure 8 . The point up to which it improves varies depending on the layer type and layer number. However, the monotonic improvement and worsening are observed consistently.

What is the effect of removing the layer completely? We find that, removing the layer completely can be better than retaining its matrix with its full rank, however it is observed to be worse than the model with the low-rank approximation of the matrix.

### C.3 Does 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} select the same layer for different tasks?

We find that the maximum improvements on different tasks come from 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} on different layers of the model. Figure 9 shows that for GPT-J on different tasks, the best-performing models across tasks have reduced matrices in different layers.

### C.4 Measuring Perplexity on PILE.

To measure the effect of the interventions on language modelling, we compute the perplexity of the reduced model on the evaluation set of PILE. The perplexity of the fixed-length GPT-J model is evaluated using a sliding window strategy over the sequence of tokens with a stride of 512 tokens. While there is an improvement in the task at hand, the model’s perplexity worsens slightly after applying 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} . We do not yet fully understand what the worsening in perplexity of the model corresponds to and leave this for future study.

### C.5 Final 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} search results

Table 3 shows the final search results of 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} for models and datasets from Table 1 . These values are obtained by reporting the optimal LASER parameters that maximize the validation accuracy. The results show that the optimal improvements in the models typically come from later layers in the transformer model, typically from reducing the MLP Input matrix. Note that ℓ = 0 \ell=0 denotes that the intervention is done on the first layer. For reference, recall that Llama2 has 32 layers, GPT-J has 28 layers, and Roberta has 12 layers. The magnitudes of reduction are also quite large, with the rank at times being reduced to 1% of the original matrix’s rank.

## Appendix D Alternate Pruning Methods

Instead of approximating weight matrices with their rank-k approximations, we tried Absolute Weight Pruning [ Frankle and Carbin, 2018 ] . Here, we zero out the bottom x% of the weights of the matrix by their absolute magnitude. The results for GPT-J on CounterFact can be seen in Figure 10 . In this case too, we find that the accuracy of the model increases with pruning later layers of the MLP. We leave further study of this phenomenon for future work.

## Appendix E Implementation Details

### E.1 Dataset Processing Details

We process each dataset described in Table 1 separately. In each case, we use 20% of the processed dataset as the validation set which we use to select the best 𝙻𝙰𝚂𝙴𝚁 {\tt LASER} hyperparameters ( τ , ρ , ℓ ) (\tau,\rho,\ell) . This validation set can be different from the validation set of the original unprocessed dataset. We use accuracy on the validation set for selecting the best hyperparameters for each LLM and a given dataset. Table 4 summarizes the size of these filtered dataset. We describe the dataset-specific processing below:

##### CounterFact.

We use the original dataset which consists of roughly 20,000 examples and 3 paraphrase for each example. This gives us 65,757 examples for the entire dataset. The set of possible labels in this QA task is open-ended. For Roberta and GPT-J LLMs, the labels are always exactly one token, while for Llama2 the labels can be multiple tokens long.

##### Hotpot.

We combine the included validation and training sets of Hotpot to increase the dataset size. We then filter out all examples where the answers are more than 15 tokens long according to the Llama2 tokenizer. We convert the original question to the prompt “<question> The answer is” (if the question ends with ? or .) or “<question>? The answer is” where the prompt variable <question> is replaced by the original question. This gives us a dataset of size 14,618. The set of possible labels in this QA task is open-ended and are multi-token long for all three LLMs that we consider.

##### Fever.

We merge the dev and test split of the original Fever dataset. We then filter out samples where with duplicate claims (inputs) but different labels (outputs). This results in a dataset of 13,086 samples, including 6,510 from the original dev set. Here there are only two possible labels: true and false. We convert each question into the prompt “Consider the following claim: <question>. Is this claim true or false. The claim is” . The prompt variable <question> is replaced by the original question.

##### Bios Gender.

We use only the dev split of the original Bias in Bios dataset. This gives us a dataset of size 39,642. The only possible labels in this QA task are two: male and female. We convert each input bio into the prompt “Consider the following text: <bio>. Is the person in this text male or female? The person is” . The prompt variable <bio> is replaced by the original bio.

##### Bios Profession.

We use only the dev split of the original Bias in Bios dataset. The goal here is to predict the profession for a given bio. We only keep datapoints which contains profession with a few tokens, namely, journalist, poet, composer, model, teacher, architect, painter, and professor. This gives us a dataset of size 19,223. The aforementioned professions compose the list of possible labels. We convert each input bio into the prompt “Consider the following text: <bio>. What is the profession of the person in this text? The profession of this person is” . The prompt variable <bio> is replaced by the original bio.

##### BigBench Epistemic Reasoning.

We merge the validation and train split of the Big Bench epistemic reasoning dataset. This gives us a dataset of size 2000. The set of possible labels here are: entailment and non-entailment which are multi-token long for all LLMs. We do not process the text.

##### Truthful QA.