mlsanityData to 7 Oct 2026

Visualising Information Flow in Word Embeddings with Diffusion Tensor Imaging

arXiv 2601.05713 ↗First posted 9 Jan 2026cs.CL cs.AI cs.LGcited by 0

Not on the map. Why?

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.

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

Builds on

36 of 70 references.

References · 36

  1. 2017Attention Is All You Needcited by 59,932
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
  3. 2013Efficient Estimation of Word Representations in Vector Spacecited by 4,346
  4. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
  5. 2020Longformer: The Long-Document Transformercited by 3,362
  6. 2018UMAP: Uniform Manifold Approximation and Projection for Dimension Reductioncited by 2,841
  7. 2019PEGASUS: Pre-training with Extracted Gap-sentences for Abstractive Summarizationcited by 1,100
  8. 2020A Primer in BERTology: What we know about how BERT workscited by 807
  9. 2023Large Language Modelscited by 774
  10. 2019How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddingscited by 770
  11. 2022Mixture-of-Experts with Expert Choice Routingcited by 611
  12. 2017BranchyNet: Fast Inference via Early Exiting from Deep Neural Networkscited by 591
24 more
  1. 2020What is the State of Neural Network Pruning?cited by 578
  2. 2019What do you learn from context? Probing for sentence structure in contextualized word representationscited by 436
  3. 2019Reducing Transformer Depth on Demand with Structured Dropoutcited by 434
  4. 2017SkipNet: Learning Dynamic Routing in Convolutional Networkscited by 328
  5. 2025Layer by Layer: Uncovering Hidden Representations in Language Modelscited by 293
  6. 2021Sustainable AI: Environmental Implications, Challenges and Opportunitiescited by 242
  7. 2022The Carbon Footprint of Machine Learning Training Will Plateau, Then Shrinkcited by 183
  8. 2022Structured Pruning Learns Compact and Accurate Modelscited by 167
  9. 2021Memorization vs. Generalization: Quantifying Data Leakage in NLP Performance Evaluationcited by 69
  10. 2020Attention Flows: Analyzing and Comparing Attention Mechanisms in Language Modelscited by 49
  11. 2016Embedding Projector: Interactive Visualization and Interpretation of Embeddingscited by 43
  12. 2024Dissecting Language Models: Machine Unlearning via Selective Pruningcited by 38
  13. 2024From Redundancy to Relevance: Information Flow in LVLMs Across Reasoning Taskscited by 37
  14. 2020Theory of gating in recurrent neural networkscited by 19
  15. 2025ToMoE: Converting Dense Large Language Models to Mixture-of-Experts through Dynamic Structural Pruningcited by 10
  16. 2023SensePOLAR: Word sense aware interpretability for pre-trained contextual word embeddingscited by 9
  17. 2020Visually Analyzing Contextualized Embeddingscited by 7
  18. 2024Computing Within Limits: An Empirical Study of Energy Consumption in ML Training and Inferencecited by 4
  19. 2025Explainable Mapper: Charting LLM Embedding Spaces Using Perturbation-Based Explanation and Verification Agentscited by 4
  20. 2025Pruning Large Language Models by Identifying and Preserving Functional Networkscited by 4
  21. 2019Multilingual, Multi-scale and Multi-layer Visualization of Intermediate Representationscited by 2
  22. 2025LayerFlow: Layer-wise Exploration of LLM Embeddings using Uncertainty-aware Interlinked Projectionscited by 2
  23. 2025The Generalization Ridge: Information Flow in Natural Language Generationcited by 2
  24. 2025Emergent effects of scaling on the functional hierarchies within large language modelscited by 1

Built on by

Not cited yet.