Interpretability of Language Models via Task Spaces (2406.06441v1)

Abstract: The usual way to interpret LLMs (LMs) is to test their performance on different benchmarks and subsequently infer their internal processes. In this paper, we present an alternative approach, concentrating on the quality of LM processing, with a focus on their language abilities. To this end, we construct 'linguistic task spaces' -- representations of an LM's language conceptualisation -- that shed light on the connections LMs draw between language phenomena. Task spaces are based on the interactions of the learning signals from different linguistic phenomena, which we assess via a method we call 'similarity probing'. To disentangle the learning signals of linguistic phenomena, we further introduce a method called 'fine-tuning via gradient differentials' (FTGD). We apply our methods to LLMs of three different scales and find that larger models generalise better to overarching general concepts for linguistic tasks, making better use of their shared structure. Further, the distributedness of linguistic processing increases with pre-training through increased parameter sharing between related linguistic tasks. The overall generalisation patterns are mostly stable throughout training and not marked by incisive stages, potentially explaining the lack of successful curriculum strategies for LMs.

Interpretability of LLMs via Task Spaces

Introduction

LLMs (LMs) have achieved remarkable prowess in generating text indistinguishable from human language, marking significant progress in NLP. However, this sophistication also brings considerable challenges in interpreting the internal mechanisms of LMs, particularly regarding the cognitive or linguistic quality of their solutions. This paper explores an alternative approach to LM interpretability by constructing `linguistic task spaces,' which elucidate how LMs conceptualize various language phenomena.

Linguistic Task Spaces and Similarity Probing

The paper introduces the concept of linguistic task spaces, inspired by former works in multi-task learning (MTL), which represent an LM's generalization behavior across linguistic tasks. The authors propose a novel method called similarity probing' to estimate linguistic similarity among tasks. This method involves selectively fine-tuning LMs on specific linguistic phenomena and evaluating the effects on other tasks. Additionally, the paper presentsfine-tuning via gradient differentials' (FTGD), a technique aimed at disentangling linguistic tasks from their natural language context.

Methods

Fine-Tuning via Gradient Differentials (FTGD)

FTGD addresses the problem of `linguistic entanglement,' where linguistic tasks are inherently interwoven within natural language data. By employing minimal pairs (sentences almost identical except for a crucial grammatical difference), FTGD estimates gradients for grammatical and ungrammatical examples, calculating their differentials (i.e., the difference between the gradients). This selective tuning method significantly reduces the number of trainable parameters while retaining a large portion of the gradient information, thus isolating specific linguistic tasks more effectively.

Similarity Probing

Similarity probing determines the similarity between linguistic tasks using two perspectives:

1. Transfer Probing: Evaluates the transfer learning effects by fine-tuning an LM on one task and measuring the performance changes in others. Positive transfer indicates high similarity, negative transfer suggests interference, and no change signifies unrelated tasks.
2. Gradient Probing: Analyzes the subspace overlap and gradient alignment between tasks. High subspace overlap and aligned gradients imply beneficial transfer, while low alignment within overlapping subspaces indicates interference.

Experiments and Results

The paper applies these methods to three LMs of different scales, pre-trained on an English Wikipedia corpus and evaluated using the BLiMP benchmark. The key experiments include:

Effectiveness of FTGD

FTGD outperforms full-gradient fine-tuning in enhancing selective linguistic tasks without disrupting the LM's broader language generation capabilities. This selectivity and effectiveness highlight FTGD's potential as an efficient fine-tuning approach.

Linguistic Task Spaces

By constructing task spaces using transfer and gradient probing, the authors reveal that LMs tend to generalize linguistic tasks in accordance with higher-level linguistic structures rather than low-level vocabulary overlaps. Larger models exhibit stronger and faster generalization within linguistic phenomena compared to smaller models.

Learning Dynamics

The development of task spaces across different pre-training stages shows remarkable stability of generalization patterns, suggesting a continuous and incremental improvement in linguistic generalization without significant shifts in learned conceptual structures. This stability contrasts with human learning paradigms characterized by distinct learning stages.

Implications and Future Developments

Theoretical Implications

The findings suggest that LMs' cognitive and linguistic capabilities develop in a more distributed manner, with increasing parameter sharing and alignment between related tasks over time. The paper highlights a potential inverse relationship between tasks' intrinsic and extrinsic dimensionality during LM training, warranting further exploration.

Practical Implications

The presented methodologies offer a robust framework for interpreting LMs and can be extended to other domains such as numerical reasoning and concept learning. Moreover, the potential for explicit linguistic hypothesis testing promises to bridge formal linguistic and computational linguistic research, offering insights into subtler structural similarities within language.

Conclusion

This paper's contributions in introducing linguistic task spaces, FTGD, and similarity probing mark a significant advancement in understanding LMs' internal linguistic processes. The authors effectively demonstrate the potential of these methods for both model interpretability and linguistic theory testing, paving the way for future research that could further elucidate the intricate mechanisms underpinning LMs' language processing capabilities.

Limitations

While the methods presented are promising, they rely heavily on synthetic data with minimal pairs, which may not entirely capture natural language diversity. The requirement of minimal pairs limits the approach's applicability to other domains without such data. Future work could explore alternative ways to define `anchors' for spanning the conceptual space and validating findings across more diverse datasets.

Acknowledgements

The authors acknowledge the support and feedback from the COLT group at UPF, computational resources from the European Research Council, and the contributions of all collaborators that enriched this research.