# Visualising Information Flow in Word Embeddings with Diffusion Tensor Imaging

arXiv: https://arxiv.org/abs/2601.05713 · HTML: https://mlsanity.com/paper/2601.05713/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 36 of its 70 references that are in the map.
- First posted on arXiv 9 Jan 2026.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [PEGASUS: Pre-training with Extracted Gap-sentences for Abstractive Summarization](https://mlsanity.com/paper/1912.08777/) (2019)
- [A Primer in BERTology: What we know about how BERT works](https://mlsanity.com/paper/2002.12327/) (2020)
- [How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings](https://mlsanity.com/paper/1909.00512/) (2019)
- [Large Language Models](https://mlsanity.com/paper/2307.05782/) (2023)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [What is the State of Neural Network Pruning?](https://mlsanity.com/paper/2003.03033/) (2020)
- [Reducing Transformer Depth on Demand with Structured Dropout](https://mlsanity.com/paper/1909.11556/) (2019)
- [What do you learn from context? Probing for sentence structure in contextualized word representations](https://mlsanity.com/paper/1905.06316/) (2019)
- [SkipNet: Learning Dynamic Routing in Convolutional Networks](https://mlsanity.com/paper/1711.09485/) (2017)
- [Layer by Layer: Uncovering Hidden Representations in Language Models](https://mlsanity.com/paper/2502.02013/) (2025)
- [Sustainable AI: Environmental Implications, Challenges and Opportunities](https://mlsanity.com/paper/2111.00364/) (2021)
- [The Carbon Footprint of Machine Learning Training Will Plateau, Then Shrink](https://mlsanity.com/paper/2204.05149/) (2022)
- [Structured Pruning Learns Compact and Accurate Models](https://mlsanity.com/paper/2204.00408/) (2022)
- [Memorization vs. Generalization: Quantifying Data Leakage in NLP Performance Evaluation](https://mlsanity.com/paper/2102.01818/) (2021)
- [Attention Flows: Analyzing and Comparing Attention Mechanisms in Language Models](https://mlsanity.com/paper/2009.07053/) (2020)
- [Embedding Projector: Interactive Visualization and Interpretation of Embeddings](https://mlsanity.com/paper/1611.05469/) (2016)
- [Dissecting Language Models: Machine Unlearning via Selective Pruning](https://mlsanity.com/paper/2403.01267/) (2024)
- [From Redundancy to Relevance: Information Flow in LVLMs Across Reasoning Tasks](https://mlsanity.com/paper/2406.06579/) (2024)
- [Theory of gating in recurrent neural networks](https://mlsanity.com/paper/2007.14823/) (2020)
- [ToMoE: Converting Dense Large Language Models to Mixture-of-Experts through Dynamic Structural Pruning](https://mlsanity.com/paper/2501.15316/) (2025)
- [SensePOLAR: Word sense aware interpretability for pre-trained contextual word embeddings](https://mlsanity.com/paper/2301.04704/) (2023)
- [Visually Analyzing Contextualized Embeddings](https://mlsanity.com/paper/2009.02554/) (2020)
- [Computing Within Limits: An Empirical Study of Energy Consumption in ML Training and Inference](https://mlsanity.com/paper/2406.14328/) (2024)

## Abstract (arXiv)

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 embedding space representations change between tokens. Tracking these changes within the layers of an LLM allows for comparing different model structures and could potentially reveal opportunities for pruning an LLM's under-utilised layers. Our results show that our visualisation method 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.