InnerThoughts: Disentangling Representations and Predictions in Large Language Models

TL;DR

This paper introduces a method to disentangle the representational and predictive capabilities of large language models by learning a separate predictor network, leading to improved performance on challenging benchmarks with less computational effort.

Contribution

It proposes a novel framework that separates representation learning from prediction in LLMs, enhancing performance without extensive fine-tuning.

Findings

Significant performance improvements on hard benchmarks

Comparable results to supervised fine-tuning

Reduced computational cost

Abstract

Large language models (LLMs) contain substantial factual knowledge which is commonly elicited by multiple-choice question-answering prompts. Internally, such models process the prompt through multiple transformer layers, building varying representations of the problem within its hidden states. Ultimately, however, only the hidden state corresponding to the final layer and token position are used to predict the answer label. In this work, we propose instead to learn a small separate neural network predictor module on a collection of training questions, that take the hidden states from all the layers at the last temporal position as input and outputs predictions. In effect, such a framework disentangles the representational abilities of LLMs from their predictive abilities. On a collection of hard benchmarks, our method achieves considerable improvements in performance, sometimes comparable to supervised fine-tuning procedures, but at a fraction of the computational cost.