InnerThoughts: Disentangling Representations and Predictions in Large Language Models
This work addresses the efficiency and accuracy of knowledge extraction in LLMs for AI researchers, though it is incremental as it builds on existing prompting techniques.
The authors tackled the problem of eliciting factual knowledge from large language models (LLMs) by proposing a method that disentangles representations from predictions, using a separate neural network predictor on hidden states from all layers, which achieved performance improvements comparable to supervised fine-tuning at lower computational cost on hard benchmarks.
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.