Visualising Information Flow in Word Embeddings with Diffusion Tensor Imaging

Thomas Fabian

REVIEW 3 major objections 4 minor 8 references

Applying diffusion tensor imaging to word embeddings produces a layer-by-layer map of information flow between tokens, and the per-layer utilisation rates it yields can single out under-used layers that are candidates for pruning.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-04 06:20 UTC pith:GWGAG6BB

load-bearing objection New DTI visualisation idea, but row-wise normalisation distorts the very quantities it claims to measure. the 3 major comments →

arxiv 2601.05713 v2 pith:GWGAG6BB submitted 2026-01-09 cs.CL cs.AIcs.LG

Visualising Information Flow in Word Embeddings with Diffusion Tensor Imaging

Thomas Fabian This is my paper

classification cs.CL cs.AIcs.LG

keywords word embeddingsdiffusion tensor imaginginterpretabilitylarge language modelsembedding visualisationstructure tensorlayer pruninganisotropy

verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

This paper proposes DONALD-D, a visualisation method that treats an LLM's hidden states for a sentence as a token-by-layer matrix and applies diffusion tensor imaging (DTI) to it. Each cell of the matrix is assigned a 2×2 structure tensor built from local gradients, whose eigenvectors and eigenvalues are rendered as diffusion ellipsoids that show the dominant direction and strength of information flow. The central claim is that these flows reveal how LLMs represent whole natural language expressions, something point-plot embedding visualisations cannot do. From the tensors the method derives a per-layer utilisation rate; the paper argues that layers with consistently low utilisation are under-used and therefore natural candidates for structured pruning. A sympathetic reader would care because the method gives an interpretable, linear-time window into where token-to-token information is actually encoded inside a transformer.

Core claim

The central discovery is that information flow in an LLM's embedding space can be measured and visualised with the same tensor machinery used to image water diffusion in the brain. After averaging hidden units and row-normalising, the gradient field of the token-layer matrix yields a structure tensor at every token-layer cell; the principal eigenvector points in the dominant flow direction and the normalised eigenvalue difference gives an anisotropy value. Applied to BERT, Longformer, GPT-2 and PEGASUS on an example sentence, DONALD-D shows distinct, model-specific flow patterns — BERT and Longformer differ despite similar encoder-only architectures, GPT-2 shows a pronounced middle band of t

What carries the argument

The central object is the 2×2 structure tensor J_{i,j} computed at every token-layer cell of the normalised hidden-state matrix M. Its components are smoothed products of the horizontal and vertical finite-difference gradients, Jxx = (∂xM)², Jxy = ∂xM ∂yM, Jyy = (∂yM)². Eigen-decomposition yields the principal flow direction v1 and anisotropy A = (λ1−λ2)/(λ1+λ2), which are drawn as oriented diffusion ellipsoids and colour-coded tiles. The same tensor components give the per-layer utilisation rate U_i = mean_j Jxx/(Jxx+Jyy), which quantifies how much of a layer's activity is token-to-token rather than perpendicular. This tensor carries the whole argument: all visualisations, anisotropy fields

Load-bearing premise

The whole pipeline rests on the assumption that after averaging the hidden units and row-normalising, the local slopes of the resulting matrix faithfully represent 'information flow' between tokens and layers — if that collapsed matrix no longer captures the embedding structure, every visualisation and pruning recommendation built on it is unsupported.

What would settle it

Take a transformer such as BERT and a moderately large corpus; compute U_i per layer with DONALD-D, then prune the three layers with the lowest U_i (by zeroing or removal) and compare task performance against pruning three randomly chosen layers. If low-U pruning does not clearly outperform random pruning, the claim that low utilisation identifies redundant layers fails. A second check: shuffle or randomly re-initialise the hidden units before averaging; if the flow visualisation barely changes, the mean collapse is discarding the very structure the method claims to show.

Watch this falsifier. Get emailed when new claim-graph text bears on it.

If this is right

Existing word-embedding visualisations compare isolated points; DONALD-D makes the context itself visible, so claims about how an LLM represents a whole expression become directly inspectable.
A consistently low utilisation rate in a layer indicates that the layer contributes little token-discriminating information, making it a concrete pruning target; corpus-level U_i values could power task-specific pruning without retraining.
Because the pipeline is model-agnostic and linear in token count, it can be applied to any transformer's hidden states, from BERT-style encoders to GPT-2-style decoders, and reveals structural differences between them.
Task-specific flow signatures — pronoun resolution in layers 3–6, metaphor detection in lower layers — offer a way to localise where linguistic phenomena are processed, potentially guiding targeted fine-tuning or intervention.
The colour map indicates orientation, not direction; since tokens are ordered, red-to-blue flow is unambiguous, letting the heatmap be read without a sign convention.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The mapping from low utilisation to redundancy is not tested in the paper; an obvious next experiment is to prune the lowest-U layers on a large corpus and compare task accuracy against random or magnitude pruning — that would either validate or falsify the pruning claim.
Because the method collapses hidden units by a mean, it cannot distinguish whether low flow reflects genuine redundancy or merely cancellation in the mean; extending DTI to the full hidden-unit axis (3D structure tensors) would test whether the visualised flows survive at unit level.
The utilisation rate depends on the input sentence; aggregated over a corpus it may correlate with known properties such as layer-wise probing results, providing a cheap, unsupervised complement to probing classifiers.
Since the gradient fields are computed on a row-normalised matrix, the method is sensitive to the normalisation choice; applying the same pipeline to attention weights or residual streams could reveal whether the observed flows are specific to hidden-state geometry or generalise to other internals.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Desk Editor's Note

New DTI visualisation idea, but row-wise normalisation distorts the very quantities it claims to measure.

read the letter

I've read the paper and the stress-test note. The note is right, and it's the thing to care about.

What's actually new: applying structure tensors / DTI to the layer-token matrix of an LLM's hidden states is a fresh idea, and the code is public. The writing is clear, the equations are standard and correct, and the O(T) complexity claim holds. The minimal-pair examples (pronoun, metaphor) are a sensible way to show the tool can produce different images.

The problem is that the central quantity — 'information flow' — is defined on a matrix that has been row-wise normalised per layer (Eq. 2). That preprocessing rescales every layer to [0,1] using its own min and max, so the layer-to-layer gradient ∂yM (Eq. 4) mixes the original vertical differences with the per-layer ranges. The utilisation rate U_i (Eq. 10) is therefore driven by the inverse of each row's range, not by how much token-specific information the layer actually carries. A layer with tiny absolute variation across tokens gets amplified and can look highly utilised; a layer with large variation can look under-utilised. The paper says normalisation makes layers 'comparable', but for the vertical gradient it does the opposite. This is a mathematical distortion built into the core computation, not just an unvalidated parameter.

The paper doesn't validate the method against any ground truth, runs no ablations (e.g., without normalisation, or with a scale-invariant gradient), and the only evidence is a few single-sentence examples. The limitations section acknowledges the mean-collapse but not this. So the pruning-suggestion claims are speculative, and the stronger 'novel insights into how LLMs represent language' claim is not supported.

That said, the paper isn't sloppy in its own terms — it's a clearly written methods paper that just overreaches. The reader's conditional verdict is reasonable; I'd say the soundness is a bit lower than a 5 given the load-bearing role of Eq. 2. But the idea is worth a referee's time, mainly so the issues are articulated and the method is repositioned as a visualisation of the normalised matrix, not a measure of true information flow.

For peer review: yes, send it. Expect heavy revision. I wouldn't cite it in my own work as a positive source, but I'd bring it to reading group to discuss what 'information flow' can and cannot mean after preprocessing.

Referee Report

3 major / 4 minor

Summary. The paper introduces DONALD-D, a method that applies diffusion tensor imaging (DTI) to word embeddings of whole sentences. The authors collapse the hidden-unit dimension of an LLM's token-layer hidden states by arithmetic mean (Eq. 1), row-wise min-max normalise the resulting L×T matrix (Eq. 2), estimate structure tensors from horizontal and vertical gradients (Eqs. 3-9), and visualise the results as colour-coded diffusion ellipsoids. A per-layer utilisation rate U_i (Eq. 10) is derived from the structure-tensor components and interpreted as the fraction of token-to-token 'information flow'. The paper presents qualitative visualisations for BERT, Longformer, GPT-2, and PEGASUS on one sentence, and two minimal-pair analyses for pronoun resolution and metaphor detection in BERT. The authors suggest that under-utilised layers, identified by low U_i, could be pruned.

Significance. If the method were validated, DONALD-D would be a novel visualisation tool connecting DTI to NLP interpretability, with potential applications in model comparison and structured pruning. The mathematical derivation of the structure tensors from the normalised matrix is standard and correctly presented, and the code is publicly released. However, the paper's central interpretive claim — that the computed gradients and utilisation rates measure 'information flow' — is not supported by any external validation, and the pruning implications are stated as an assumption rather than demonstrated. The method is deterministic and has only one free parameter (the smoothing kernel), which is a strength in terms of reproducibility.

major comments (3)

[§2.1, Eq. (2)] Row-wise min-max normalisation is load-bearing for the layer-to-layer gradient and utilisation rate. Because each row is rescaled by its own minimum and range, the vertical gradient ∂yM in Eq. (4) is not a monotone transform of the original layer-to-layer differences: ∂yM'_i,j = (M_{i+1,j}-m_{i+1})/(2s_{i+1}) - (M_{i-1,j}-m_{i-1})/(2s_{i-1}), where s_i = max_j M_i,j - min_j M_i,j. Thus Jyy (Eq. 8) and U_i (Eq. 10) depend on the per-layer ranges s_i, which are arbitrary scaling factors unrelated to the embedding information content. A layer with tiny token-wise variation is amplified to the full [0,1] range, making it appear utilised along the token direction even if its absolute variation is negligible; conversely, a layer with large absolute token differences is suppressed. The paper's claim that this normalisation makes layers 'comparable' (Sec. 2.1) is therefore incorrect for the laye
[§3 and §4] The central construct 'information flow' is never validated against an external measure. The visualisations and U_i are computed from the arithmetic mean of hidden states (Eq. 1) and normalised gradients; no evidence is given that these quantities track token/layer information in the LLM's representations. The interpretive claims in §3, such as 'this indicates that the semantic token-to-token differences increase when the expression is used metaphorically', are not backed by any control, ground-truth comparison, or quantitative evaluation. I would expect at least one validation experiment: for example, correlate U_i with known layer functions from probing classifiers, or show that the visualisation changes as expected when the input is randomly perturbed, or compare against attention/entropy measures. Without such validation, the paper demonstrates only that DTI can be applied to a norma
[§2.2 and §4] The pruning conclusion is not experimentally supported. Eq. (10) is interpreted as indicating that a layer with low U_i 'might be redundant' in the embedding space, and the Discussion states that it is 'reasonable to assume that removing under-utilised layers from a model should have little impact on its performance'. This is a testable hypothesis, yet no pruning experiment is performed, nor is U_i compared with established pruning metrics or actual performance drops. Since the identification of pruning opportunities is a stated application of DONALD-D, a pilot pruning experiment on a small model (e.g., BERT-tiny) is needed to justify the claim that utilisation rates identify removable layers.

minor comments (4)

[§2.2] In the sentence describing Gaussian smoothing, '∂xM and ∂xM' should read '∂xM and ∂yM'.
[§2.4] The colourmap is described as 'π-periodic' and then as assigning the same colour to θ and θ+180°. If angles are in degrees, the period is 180°, not π; if in radians, the example should be θ+π. Please clarify consistently.
[Table 1] The table lists layers from 24 (top) to 1 (bottom) for some models and 16 to 1 for PEGASUS, but the direction of the y-axis in Figure 2 is not explicitly stated. It would help to state whether layer 1 is the first transformer block or the last, and to align the table with the figures.
[§3] The sentence 'BERT utilises layers one to eight the most' is imprecise given Table 1, where layer 8 has U=38.83% while layer 7 has 59.85%. Consider describing the profile more carefully, e.g., 'layers 1-7 have high utilisation, with a drop at layer 8'.

Circularity Check

0 steps flagged

No circularity: DONALD-D is a descriptive visualisation pipeline with no fitted predictions, no imported uniqueness theorems, and no self-citation chain.

full rationale

Walking the claimed derivation chain: hidden states E are collapsed by Eq. 1, row-wise min-max normalised by Eq. 2, converted to finite-difference gradients by Eqs. 3–4, assembled into structure tensors by Eqs. 5–9, and then used to define anisotropy and utilisation by Eqs. 10–12. Every plotted or tabulated quantity is defined by construction from the input embeddings; none is fitted to an external target and then reported as a prediction. The abstract claims that the method 'permits novel insights', not that an independently defined quantity was derived and confirmed. The closest candidate for circularity is Eq. 2, where row-wise normalisation changes the scale of the layer-to-layer gradient and therefore affects U_i; however, the paper does not claim U_i is an independent measure of layer utilisation that the method then predicts. This is a modelling-validity concern about whether U_i tracks embedding geometry, not a circularity concern. The Limitations section explicitly concedes the simplification: 'DONALD-D collapses the hidden units of each layer and only considers the arithmetic mean... potentially obscures changes in the hidden units', and states 'The results described in this work are exemplary'. There are no self-citations, no author-imported uniqueness theorems, no ansatz smuggled in by citation, and no fitted input renamed as a prediction. DTI is borrowed from neuroimaging as an explicit analogy, and the layer-semantics interpretation cites the external, non-overlapping work of Jawahar et al. (2019). The method is self-contained; its weakness is interpretive validity, not circularity.

Axiom & Free-Parameter Ledger

1 free parameters · 4 axioms · 0 invented entities

No fitted constants or newly postulated physical entities. The key assumptions are the averaging/normalisation choices and the interpretation of gradient anisotropy as information flow, plus the untested pruning assumption. The smoothing width is an unspecified free parameter.

free parameters (1)

Gaussian smoothing kernel width (sigma) = not specified
Used in Section 2.2 to smooth derivative fields; the value is not given, and it affects the structure tensors, anisotropy, and utilisation rates.

axioms (4)

standard math DTI/structure-tensor formalism (Bigun & Granlund, Weickert) is valid for discrete 2D matrices.
Assumed without proof; standard in image processing and neuroimaging.
domain assumption The arithmetic mean over hidden units (Eq. 1) preserves the information-flow structure needed for the visualisation.
Load-bearing; no validation that mean activations capture token-to-token and layer-to-layer information flow.
ad hoc to paper Row-wise min-max normalisation (Eq. 2) makes layers comparable without removing token-wise variation.
Design choice unique to this tool; it changes Jxx/Jyy ratios and hence the utilisation rates.
ad hoc to paper A low utilisation rate Ui implies that a layer is redundant and can be pruned (Eq. 10 and Discussion).
The paper states this is a 'reasonable assumption' but does not test pruning.

pith-pipeline@v1.3.0-alltime-deepseek · 11259 in / 9600 out tokens · 100889 ms · 2026-08-04T06:20:52.088491+00:00 · methodology

0 comments

read the original 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 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.

Figures

Figures reproduced from arXiv: 2601.05713 by Thomas Fabian.

**Figure 2.** Figure 2: Information flow in an exemplary natural language expression in BERT ( [PITH_FULL_IMAGE:figures/full_fig_p006_2.png] view at source ↗

**Figure 3.** Figure 3: Minimal pair analysis for pronoun resolution. [PITH_FULL_IMAGE:figures/full_fig_p007_3.png] view at source ↗

**Figure 4.** Figure 4: Minimal pair analysis for metaphor detection. [PITH_FULL_IMAGE:figures/full_fig_p007_4.png] view at source ↗

discussion (0)

Reference graph

Works this paper leans on

8 extracted references · 5 linked inside Pith

[6]

Place: Cambridge, MA Publisher: MIT Press

A Primer in BERTology: What We Know About How BERT Works.Transactions of the As- sociation for Computational Linguistics, 8:842–866. Place: Cambridge, MA Publisher: MIT Press. Cynthia Rudin. 2019. Stop explaining black box ma- chine learning models for high stakes decisions and use interpretable models instead.Nature Machine Intelligence, 1(5):206–215. Da...

Pith/arXiv arXiv 2019
[1513]

Publisher: John Wiley & Sons, Ltd. Carole-Jean Wu, Ramya Raghavendra, Udit Gupta, Bilge Acun, Newsha Ardalani, Kiwan Maeng, Glo- ria Chang, Fiona Aga, Jinshi Huang, Charles Bai, Michael Gschwind, Anurag Gupta, Myle Ott, Anas- tasia Melnikov, Salvatore Candido, David Brooks, Geeta Chauhan, Benjamin Lee, Hsien-Hsin Lee, and 6 others. 2022. Sustainable AI: E...

Pith/arXiv arXiv 2022
[1994]

Publisher: Elsevier

MR diffusion tensor spectroscopy and imag- ing.Biophysical Journal, 66(1):259–267. Publisher: Elsevier. Iz Beltagy, Matthew E. Peters, and Arman Cohan
[2019]

Association for Computational Linguistics

What Does BERT Learn about the Structure of Language? InProceedings of the 57th Annual Meet- ing of the Association for Computational Linguistics, pages 3651–3657, Florence, Italy. Association for Computational Linguistics. Denis LeBihan, Jean-François Mangin, Cyril Poupon, Chris A. Clark, Sabina Pappata, Nicolas Molko, and Hughes Chabriat. 2001. Diffusio...

2001
[2020]

ArXiv:2004.05150 [cs]

Longformer: The Long-Document Trans- former.arXiv preprint. ArXiv:2004.05150 [cs]. Emily M. Bender and Alexander Koller. 2020. Climbing towards NLU: On Meaning, Form, and Understand- ing in the Age of Data. InProceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 5185–5198, Online. Association for Computational Ling...

Pith/arXiv arXiv 2004
[2022]

Lorenzo Pini, Michela Pievani, Martina Bocchetta, Daniele Altomare, Paolo Bosco, Enrica Cavedo, Samantha Galluzzi, Moira Marizzoni, and Gio- vanni B

The Carbon Footprint of Machine Learn- ing Training Will Plateau, Then Shrink.Computer, 55(7):18–28. Lorenzo Pini, Michela Pievani, Martina Bocchetta, Daniele Altomare, Paolo Bosco, Enrica Cavedo, Samantha Galluzzi, Moira Marizzoni, and Gio- vanni B. Frisoni. 2016. Brain atrophy in Alzheimer’s Disease and aging.Brain Imaging and Aging, 30:25– 48. Alec Rad...

2016
[2023]

InFindings of the As- sociation for Computational Linguistics: ACL 2023, pages 10880–10895, Toronto, Canada

Structured Pruning for Efficient Generative Pre-trained Language Models. InFindings of the As- sociation for Computational Linguistics: ACL 2023, pages 10880–10895, Toronto, Canada. Association for Computational Linguistics. Surat Teerapittayanon, Bradley McDanel, and Hsiang Tsung Kung. 2016. BranchyNet: Fast infer- ence via early exiting from deep neural...

Pith/arXiv arXiv 2023
[2024]

ArXiv:2406.14328 [cs]

Computing Within Limits: An Empirical Study of Energy Consumption in ML Training and Inference.arXiv preprint. ArXiv:2406.14328 [cs]. Leland McInnes, John Healy, and James Melville. 2020. UMAP: Uniform Manifold Approximation and Pro- jection for Dimension Reduction.arXiv preprint. ArXiv:1802.03426 [stat]. Tomas Mikolov, Kai Chen, Greg Corrado, and Jef- fr...

Pith/arXiv arXiv 2020

[1] [6]

Place: Cambridge, MA Publisher: MIT Press

A Primer in BERTology: What We Know About How BERT Works.Transactions of the As- sociation for Computational Linguistics, 8:842–866. Place: Cambridge, MA Publisher: MIT Press. Cynthia Rudin. 2019. Stop explaining black box ma- chine learning models for high stakes decisions and use interpretable models instead.Nature Machine Intelligence, 1(5):206–215. Da...

Pith/arXiv arXiv 2019

[2] [1513]

Publisher: John Wiley & Sons, Ltd. Carole-Jean Wu, Ramya Raghavendra, Udit Gupta, Bilge Acun, Newsha Ardalani, Kiwan Maeng, Glo- ria Chang, Fiona Aga, Jinshi Huang, Charles Bai, Michael Gschwind, Anurag Gupta, Myle Ott, Anas- tasia Melnikov, Salvatore Candido, David Brooks, Geeta Chauhan, Benjamin Lee, Hsien-Hsin Lee, and 6 others. 2022. Sustainable AI: E...

Pith/arXiv arXiv 2022

[3] [1994]

Publisher: Elsevier

MR diffusion tensor spectroscopy and imag- ing.Biophysical Journal, 66(1):259–267. Publisher: Elsevier. Iz Beltagy, Matthew E. Peters, and Arman Cohan

[4] [2019]

Association for Computational Linguistics

What Does BERT Learn about the Structure of Language? InProceedings of the 57th Annual Meet- ing of the Association for Computational Linguistics, pages 3651–3657, Florence, Italy. Association for Computational Linguistics. Denis LeBihan, Jean-François Mangin, Cyril Poupon, Chris A. Clark, Sabina Pappata, Nicolas Molko, and Hughes Chabriat. 2001. Diffusio...

2001

[5] [2020]

ArXiv:2004.05150 [cs]

Longformer: The Long-Document Trans- former.arXiv preprint. ArXiv:2004.05150 [cs]. Emily M. Bender and Alexander Koller. 2020. Climbing towards NLU: On Meaning, Form, and Understand- ing in the Age of Data. InProceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 5185–5198, Online. Association for Computational Ling...

Pith/arXiv arXiv 2004

[6] [2022]

Lorenzo Pini, Michela Pievani, Martina Bocchetta, Daniele Altomare, Paolo Bosco, Enrica Cavedo, Samantha Galluzzi, Moira Marizzoni, and Gio- vanni B

The Carbon Footprint of Machine Learn- ing Training Will Plateau, Then Shrink.Computer, 55(7):18–28. Lorenzo Pini, Michela Pievani, Martina Bocchetta, Daniele Altomare, Paolo Bosco, Enrica Cavedo, Samantha Galluzzi, Moira Marizzoni, and Gio- vanni B. Frisoni. 2016. Brain atrophy in Alzheimer’s Disease and aging.Brain Imaging and Aging, 30:25– 48. Alec Rad...

2016

[7] [2023]

InFindings of the As- sociation for Computational Linguistics: ACL 2023, pages 10880–10895, Toronto, Canada

Structured Pruning for Efficient Generative Pre-trained Language Models. InFindings of the As- sociation for Computational Linguistics: ACL 2023, pages 10880–10895, Toronto, Canada. Association for Computational Linguistics. Surat Teerapittayanon, Bradley McDanel, and Hsiang Tsung Kung. 2016. BranchyNet: Fast infer- ence via early exiting from deep neural...

Pith/arXiv arXiv 2023

[8] [2024]

ArXiv:2406.14328 [cs]

Computing Within Limits: An Empirical Study of Energy Consumption in ML Training and Inference.arXiv preprint. ArXiv:2406.14328 [cs]. Leland McInnes, John Healy, and James Melville. 2020. UMAP: Uniform Manifold Approximation and Pro- jection for Dimension Reduction.arXiv preprint. ArXiv:1802.03426 [stat]. Tomas Mikolov, Kai Chen, Greg Corrado, and Jef- fr...

Pith/arXiv arXiv 2020