Visualising Information Flow in Word Embeddings with Diffusion Tensor Imaging

TL;DR

This paper introduces a novel diffusion tensor imaging approach to visualize and analyze information flow in word embeddings within large language models, enhancing interpretability of how models process natural language expressions.

Contribution

The paper presents a new DTI-based method for analyzing information flow in LLMs, extending beyond single words to entire expressions and enabling comparison of model structures.

Findings

DTI reveals information flow between word embeddings

Tracking flows within layers shows opportunities for pruning

Differences in flows for pronoun resolution and metaphor detection

Abstract

Understanding how large language models (LLMs) represent natural language is a central challenge in natural language processing (NLP) research. Many existing methods extract word embeddings from an LLM, visualise the embedding space via point-plots, and compare the relative positions of certain words. However, this approach only considers single words and not whole natural language expressions, thus disregards the context in which a word is used. Here we present a novel tool for analysing and visualising information flow in natural language expressions by applying diffusion tensor imaging (DTI) to word embeddings. We find that DTI reveals how information flows between word embeddings. Tracking information flows within the layers of an LLM allows for comparing different model structures and revealing opportunities for pruning an LLM's under-utilised layers. Furthermore, our model reveals differences in information flows for tasks like pronoun resolution and metaphor detection. Our results show that our model permits novel insights into how LLMs represent actual natural language expressions, extending the comparison of isolated word embeddings and improving the interpretability of NLP models.