Visualising Information Flow in Word Embeddings with Diffusion Tensor Imaging

TL;DR

This work introduces DONALD-D, a diffusion tensor imaging–based method to visualize information flow in word embeddings across natural language expressions. By constructing a matrix $M \in \mathbb{R}^{L\times T}$ through averaging hidden units and applying gradient-based diffusion-tensor analysis, the method yields diffusion ellipsoids that encode the dominant flow direction and anisotropy per token-layer cell, with a linear complexity in the expression length $O(T)$. The authors demonstrate the approach on multiple LLMs (e.g., BERT, Longformer, GPT-2, PEGASUS), revealing model-specific diffusion patterns and context-dependent shifts in information flow for pronoun resolution and metaphor detection, and they discuss using layer-utilisation metrics to inform pruning strategies (DUCK). Overall, DONALD-D provides an interpretable, context-aware extension to embedding visualisations, enabling new linguistic analyses and practical avenues for model simplification without retraining.

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.

Figures (4)

• Figure 1: Isotropic (left) and anisotropic (right) diffusion. For isotropic diffusion, the eigenvalues $\lambda_1$ and $\lambda_2$ are equal. The degree of anisotropy increases with the difference between the eigenvalues as a larger fraction of directional strength runs along the principal eigenvector $v_1$.
• Figure 2: Information flow in an exemplary natural language expression in BERT devlin_bert_2019, Longformer beltagy_longformer_2020, GPT-2 radford_language_2019, and PEGASUS zhang_pegasus_2020. Each model tokenises and provides word embeddings for the expression: This plot visualises information flow in word embeddings via diffusion ellipsoids. You can make this plot as large as you want by adding more and more sentences. The x-axis shows the tokens, the number of which can vary between models due to different tokenisation processes. The y-axis shows the layers of each model. The ellipsoids within each tile encode the direction of information diffusion. Red tiles indicate token-to-token information flow. Blue tiles indicate no information flow along the token sequence.
• Figure 3: Minimal pair analysis for pronoun resolution. The word embeddings are taken from BERT devlin_bert_2019. Each column corresponds to a token in the expression. The highlighted columns correspond to the bold target words, he and it, respectively.
• Figure 4: Minimal pair analysis for metaphor detection. The word embeddings are taken from BERT devlin_bert_2019. Each column corresponds to a token in the expression, highlighted columns correspond to the bold target phrase he kicked the bucket.