Visualising Information Flow in Word Embeddings with Diffusion Tensor Imaging
This work addresses the challenge of interpreting NLP models for researchers, offering a novel visualization method that goes beyond isolated word comparisons, though it appears incremental as it builds on existing embedding analysis techniques.
The authors tackled the problem of understanding how large language models represent natural language by developing a tool that applies diffusion tensor imaging to word embeddings to visualize information flow in natural language expressions, revealing insights into model structures and tasks like pronoun resolution and metaphor detection.
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.