Pith. sign in

REVIEW 5 major objections 4 minor 2 cited by

Transformers as Unrolled Inference in Probabilistic Laplacian Eigenmaps: An Interpretation and Potential Improvements

T0 review · 5 major / 4 minor · reviewed 2026-08-06 · deepseek-v4-flash

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

Pith's one-line read Transformer blocks can be read as unrolled inference in a probabilistic Laplacian Eigenmaps model, so the attention matrix should really be a graph Laplacian $A-I$, and swapping it in improves validation performance and training speed.

desk verdict The A-I attention tweak is worth a serious test, but the paper's derivation has a load-bearing flaw: with row-softmax, I-A is non-symmetric, so the A-I update is not the ELBO gradient. read the letter →

arxiv 2507.21040 v1 pith:CKP6MGX2 submitted 2025-07-28 cs.LG stat.ML

Aditya Ravuri , Neil D. Lawrence This is my paper
classification cs.LGstat.ML
keywords transformersprobabilisticLaplacianEigenmapsunrolledinferencegraphdiffusionattentionmechanismvariationaldimensionalityreductionsmoothing
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

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

The reading

This paper tries to show that a transformer block is not just analogous to but actually an unrolled optimisation step for a probabilistic Laplacian Eigenmaps model, a latent-variable dimensionality-reduction model. On that reading, the softmax attention matrix should be treated as a graph adjacency matrix, and the correct update inside attention is a graph diffusion step, i.e. $A-I$ rather than $A$. The authors derive this form from a variational lower bound and then test the architectural consequence: subtracting the identity from the attention matrix improves validation performance on a small language model and on a vision transformer, and speeds up GPT-2 training. The reason to care is that the derivation turns a heuristic attention design into a concrete inference prediction with a simple, testable modification.

What carries the argument

The load-bearing object is the variational Laplacian Eigenmaps objective $KL(q(\Gamma,Z)\|p(\Gamma,Z)) \propto \mathrm{tr}(\tilde L(XX^T+\beta I)) - \log\det(XX^T+\beta I)$, with the variational constraint $\tilde L = I - \sigma(\kappa ZZ^T - M)$ treated as an observed random variable (stop-gradient). Gradient descent on the first term produces the update $X \leftarrow X + 2\eta(\sigma(\kappa XX^T - \beta I - M)-I)X$, whose parenthesised factor is the graph Laplacian $A-I$; the second term plus a uniform-matrix prior yields the linear feed-forward update and LayerNorm projection. This $A-I$ identity is what carries the argument: it turns the attention matrix into a diffusion operator and is precisely the modification tested in the experiments.

What would settle it

Compute or numerically approximate the full gradient of the variational objective including the derivative of $\sigma(\kappa XX^T-M)$ with respect to $X$ (without stop-gradient) and check whether the resulting update is still proportional to $(A-I)X$ plus only negligible corrections. If the extra terms are not negligible, the paper's central interpretive claim fails at its stated assumption.

Watch

Extended reading notes

Core claim

At initialisation, with weights set to diagonal matrices, a transformer block performs gradient descent on the KL divergence between a variational Wishart distribution whose constraint is a soft nearest-neighbour graph Laplacian $\tilde L = I - \sigma(\kappa ZZ^T - M)$ and a model Wishart distribution over the same precision matrix. Because the derivative of the data term is $d\,\mathrm{tr}(\tilde L(XX^T+\beta I))/dX = 2\tilde L X = 2(A-I)X$, the attention operation should be the negative graph Laplacian $A-I$, i.e. a Laplacian-smoothing or graph-diffusion step, not the raw attention matrix. The paper then shows empirically that replacing the attention matrix by $A-I$ improves validation performance on Shakespeare-character and downsampled ImageNet transformers and makes GPT-2 training converge faster.

Load-bearing premise

The derivation treats the softmax attention matrix as a fixed, observed object when computing gradients, so the dependence of the attention matrix on the latent vectors $X$ is ignored; if that dependence were included, the update would acquire extra terms and the clean $A-I$ form would no longer follow.

Editorial extensions

If this is right

  • If transformer blocks are unrolled Laplacian Eigenmaps inference, then at initialisation a randomly initialised transformer performs linear dimensionality reduction and clustering of token or patch representations, as shown on MNIST.
  • The attention matrix should be read as an adjacency matrix, and graph diffusion ($A-I$) is a principled replacement for raw attention; the paper reports consistent validation gains on language and vision.
  • Training a GPT-2-scale model with the $A-I$ modification converges slightly faster than the vanilla attention run.
  • The weight matrices in a transformer can be interpreted as rotation and learning-rate (step-size) parameters, giving a concrete 'learn to optimise' story for the architecture.
  • The same derivation can be extended to kernelised dimensionality-reduction objectives, which the authors suggest as future work; if that holds, further non-linear graph-diffusion updates may improve low-dimensional models.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • If the $A-I$ result is robust, it suggests that part of attention's value is implicit Laplacian smoothing, so other graph-regularisation terms (degree normalisation, higher-order diffusion) may yield further gains; this is not tested in the paper.
  • The stop-gradient assumption is the likely breaking point: relaxing it would couple attention to the update and may change the architecture, so a natural extension is to train with the full gradient and compare.
  • The derivation may transfer to other attention-based architectures (e.g. cross-attention, multi-head with product-of-experts) even though the paper only analyses single-head self-attention at initialisation.
  • The MNIST clustering result suggests a possible diagnostic: check whether the $A-I$ modification tightens class-cluster structure in deeper layers; the paper does not report this.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

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

Referee Report

5 major / 4 minor

Summary. The paper proposes a probabilistic interpretation of transformer blocks as unrolled inference steps in a probabilistic Laplacian Eigenmaps model from the ProbDR framework. The authors argue that, at initialization, a transformer block corresponds to gradient descent on a variational lower bound, and that this derivation yields a graph Laplacian term A-I in place of the standard attention matrix A. They test the A-I modification on a language model and a vision transformer, reporting improved validation performance and faster convergence. The manuscript also includes a qualitative experiment showing that a randomly initialized transformer clusters MNIST digits in its latent space.

Significance. If the derivation were correct, the paper would offer a novel probabilistic grounding for the transformer architecture and a simple, broadly applicable architectural modification with empirical support. The paper is transparent about two key approximations (the stop-grad treatment of the variational constraint and the replacement of the regularizer gradient by a linear term), which is commendable. The empirical finding that subtracting the identity from the attention matrix improves performance is interesting even independently of the theory, and the MNIST clustering experiment provides a concrete, falsifiable illustration. However, the central derivation contains a load-bearing algebraic error regarding the gradient of a non-symmetric quadratic form, and the experiments lack statistical rigor (no seeds, no error bars). The paper does not provide code or machine-checked proofs, so the empirical claims rest on the reported figures alone.

major comments (5)
  1. [Section 3, derivative of Ldata] The equation dL/dX = 2 L X is incorrect for the non-symmetric matrix L = I - A. Because A = softmax(κ Z Z^T - M) is applied row-wise, A is generally not symmetric, so L is not symmetric. For a constant non-symmetric L, the derivative of tr(L X X^T) with respect to X is (L + L^T)X, not 2 L X. Consequently, the update X <- X + 2η(A - I)X is not the gradient-descent step of the stated ELBO, even under the paper's explicit stop-grad treatment. The correct update would involve (A + A^T)/2 - I, which is not what is implemented or tested in the experiments. This error undermines the central claim that the A-I form follows from unrolled inference.
  2. [Section 3, variational constraint and Wishart validity] The matrix L = I - A with row-stochastic softmax A is not a graph Laplacian in the sense required by the ProbDR model: it is non-symmetric and generally not positive semidefinite. A Wishart distribution W(Γ | L, d) is defined only for a positive definite scale matrix, so the variational distribution q(Γ, Z) = W(Γ | L(Z), d) is not a valid distribution for arbitrary Z. The paper should either symmetrize the Laplacian (e.g., L_sym = I - (A + A^T)/2) or justify the use of a non-PSD scale matrix. This is load-bearing because the entire ELBO simplification depends on the Wishart form.
  3. [Section 3, stop-grad on the variational constraint] The paper states, 'we treat the variational constraint as an observed random variable, and hence do not account for gradient updates to terms leading from the variational constraint.' This is a significant and nonstandard assumption. If the dependence of A on X were included in the gradient, the update would contain additional terms and the clean A-I form would not follow. The SimSiam analogy is not a sufficient justification, since stop-grad in SimSiam is an algorithmic design choice rather than a consequence of the probabilistic model. The assumption should be empirically tested, for example by comparing the proposed update with the full-gradient update on a small task.
  4. [Section 3, approximation of the regularizer gradient] The replacement dLreg/dX = 2(XX^T + βI)^{-1}X ≈ 2/(dq + β) X is stated without derivation or numerical verification. This approximation is used to argue that the feed-forward block corresponds to a gradient step, but the error in this approximation could be large, especially at initialization when XX^T may not be close to a scaled identity. The paper should either justify this approximation analytically (e.g., under the assumed initialization) or demonstrate empirically that it does not affect the qualitative conclusions.
  5. [Section 4, experimental methodology] The empirical claims rest on single runs with no reported seeds, error bars, or statistical significance. Section 4.2 says the model was 'run multiple times' but does not report the number of runs or the variance; Figure 3 appears to show a single comparison between two runs. Given the small performance differences, the claim that the A-I modification 'consistently' improves performance is not supported. The authors should report multiple seeds, mean and variance, and ideally a significance test for the language and vision experiments.
minor comments (4)
  1. [Section 4.1, Figure 1] The claim that 'transformer blocks cluster points in the latent space' is supported only by a qualitative plot of the first two latent dimensions. The paper should provide a quantitative clustering metric (e.g., silhouette score or purity) to make the dimensionality-reduction claim more precise.
  2. [Section 4.1, initialization details] The sentence 'The query and key matrices were set up such that the attention matrix, pre-normalisation, has a diagonal equal to κ' is ambiguous: the softmax output is non-negative, but the pre-softmax scores have diagonal κ. Please clarify whether κ refers to the pre-softmax or post-softmax diagonal.
  3. [Section 3, notation] The notation is inconsistent: sometimes the graph Laplacian is written as L(Y) based on data Y, and later as L(Z) based on latent Z. Please use distinct symbols or clarify the dependence to avoid confusion.
  4. [General, color-dependent text] The manuscript refers to 'the element highlighted in red' in equations, but the color is not visible in a black-and-white printout. Please use a symbol or footnote instead.

Circularity Check

1 steps flagged · score 6.0 of 10

The claimed derivation of the A−I graph-diffusion step is self-definitional: the variational graph Laplacian is defined as I−A, so the A−I update is an input assumption read back out as a result, while the empirical benchmarks remain independent evidence.

  1. self definitional [Section 3, definition of the variational graph Laplacian and the gradient update]
    "The graph Laplacian is computed as L = I - A(Z) = I - σ(κZZ^T - M) where σ is the softmax function, applied row-wise (so that the row sums of the input matrix all equal one). ... dL/dX = 2 L X = 2(A - I), and a gradient descent update ... X ← X + 2η(σ(κXX^T - βI - M) - I)X."

    The central interpretive conclusion — that a graph Laplacian term (A−I) 'arises from our arguments' — holds only because the variational constraint was defined with L = I − A, where A is the softmax attention matrix. Substituting that definition into the objective yields the A−I update immediately; no step produces A−I from the ProbDR model alone. The attention matrix is inserted into the variational constraint as a modeling choice, so the proposed architectural modification is equivalent to that choice by construction, not a first-principles prediction. The empirical improvements are independent evidence, but the derivation's claim to explain why A−I works restates its own ansatz.

full rationale

The paper is not wholly circular: the language and vision experiments (Shakespeare, downsampled ImageNet, GPT-2) provide external evidence that replacing attention A with A−I helps, and those results stand independently of the derivation. The self-citations to Ravuri et al. (2023) and Ravuri & Lawrence (2024) are not themselves load-bearing in a circular way; the ProbDR model is a starting point, not a conclusion, and the closed-form eigenvector claims rest on standard linear algebra. However, the central interpretive claim is circular by construction. The variational graph Laplacian is defined as I − A with A the row-wise softmax attention matrix, and the paper then reports that a graph Laplacian (equivalently A−I) 'arises from our arguments.' This is an assumption unfolded into a gradient update, not a derived prediction. For completeness, note a separate correctness problem: for a non-symmetric row-stochastic A, the derivative of tr((I−A)XX^T) is (2I−A−A^T)X, not 2(I−A)X, so even under the paper's stop-grad treatment of the variational constraint the displayed gradient update is algebraically wrong (and the sign of L is also flipped between L=I−A and dL/dX=2(A−I)). That is a mathematical error rather than a circularity, but it reinforces that the A−I modification is not an inevitable consequence of the stated ELBO. On balance, the interpretive derivation reduces to its own definition, giving a partial circularity score of 6, while the empirical component keeps the paper from being entirely vacuous.

Assumptions & free parameters 5 free parameters · 5 assumptions · 0 invented entities

The interpretive claim rests on several modeling choices: the ProbDR model itself, the softmax-as-adjacency proxy, the stop-grad approximation, and a crude approximation for the regularizer gradient. The empirical claim about A-I is supported by experiments, but the experiments are preliminary.

free parameters (5)
  • kappa (κ) = 30
    Set based on the clustering empirically observed in the resulting graph Laplacian's eigenvectors (Section 4.1).
  • beta (β) = not specified
    Inverse regularizer in the PLE model; appears in the derivation but no value is given for the experiments.
  • eta (η) = 0.4
    Step size; chosen with a max of 0.5 to keep update magnitudes bounded (Section 4.1).
  • dq (latent dimension) = 128
    Chosen for the MNIST experiment with no stated justification (Section 4.1).
  • M (mask matrix) = not specified
    Mask in the softmax adjacency proxy; no values given for the experiments.
assumptions (5)
  • domain assumption Probabilistic Laplacian Eigenmaps model with Wishart prior, from Ravuri et al. (2023)
    The paper builds on this model without re-deriving it (Section 2).
  • ad hoc to paper Softmax row-wise output is a soft proxy for the nearest-neighbour adjacency matrix
    Section 3: 'A tilde, we argue, is a soft proxy to the true nearest neighbour adjacency matrix'; this is a modeling choice.
  • ad hoc to paper Stop-grad on the variational constraint, treating the attention matrix as constant
    Section 3: 'we treat the variational constraint as an observed random variable, and hence do not account for gradient updates to terms leading from the variational constraint'; this ignores the dependence of the attention matrix on X.
  • ad hoc to paper Approximation dLreg/dX ≈ 2/(dq+β)X
    Section 3: 'we simply approximate dLreg/dX = 2(XXT + beta I)^{-1}X ≈ 2/(dq+beta)X'; no justification given.
  • domain assumption Analysis restricted to initialization with diagonal weights, single head, and no ReLU
    Section 3: 'we analyse the transformer at initialisation... and consider transformers with single heads... we ignore the ReLU activation'.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Transformers as Unrolled Inference in Probabilistic Laplacian Eigenmaps: An Interpretation and Potential Improvements." pith.science (2026). https://pith.science/paper/CKP6MGX2

@misc{pith2026250721040,
  author       = {Pith},
  title        = {Pith review of: Transformers as Unrolled Inference in Probabilistic Laplacian Eigenmaps: An Interpretation and Potential Improvements},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/CKP6MGX2}},
  note         = {Machine review of arXiv:2507.21040}
}
read the original abstract

We propose a probabilistic interpretation of transformers as unrolled inference steps assuming a probabilistic Laplacian Eigenmaps model from the ProbDR framework. Our derivation shows that at initialisation, transformers perform "linear" dimensionality reduction. We also show that within the transformer block, a graph Laplacian term arises from our arguments, rather than an attention matrix (which we interpret as an adjacency matrix). We demonstrate that simply subtracting the identity from the attention matrix (and thereby taking a graph diffusion step) improves validation performance on a language model and a simple vision transformer.

Figures

Figures reproduced from arXiv: 2507.21040 by the authors.

Figure 1
Figure 1. The first two latent dimensions corresponding to flattened MNIST images after a random initial [PITH_FULL_IMAGE:figures/full_fig_p005_1.png] view at source ↗
Figure 2
Figure 2. Left: validation losses on the Shakespeare dataset and right: validation accuracies on a down￾sampled Imagenet dataset, showing that Laplacian smoothing achieves a better performance in both cases. 4.3 GPT-2 converges faster with Graph Diffusion In fig. 3, we show the difference in training losses between a run without graph diffusion, and a run with our modification. We use the same pre-training strategy laid out i… view at source ↗
Figure 3
Figure 3. Visualisation of the difference in training losses with and without graph diffusion. A positive [PITH_FULL_IMAGE:figures/full_fig_p006_3.png] view at source ↗

Discussion (0). Sign in to comment.

Forward citations

Cited by 2 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. Revisiting Transformer Layer Parameterization Through Causal Energy Minimization

    cs.LG 2026-05 unverdicted novelty 6.0 of 10

    CEM recasts Transformer layers as energy minimization steps, enabling constrained parameterizations like weight sharing and low-rank interactions that match standard baselines in 100M-scale language modeling.

  2. Uncovering and Shaping the Latent Representation of 3D Scene Topology in Vision-Language Models

    cs.CV 2026-05 unverdicted novelty 6.0 of 10

    VLMs possess a latent 3D scene topology subspace corresponding to Laplacian eigenmaps that can be causally shaped via Dirichlet energy regularization to improve spatial task performance by up to 12.1%.

Reference graph

Works this paper leans on

15 extracted references · 5 canonical work pages · cited by 2 Pith papers

  1. [3]

    Learning to Optimize: A Primer and A Benchmark

    URLhttps://arxiv.org/abs/2103.12828. Xinlei Chen and Kaiming He. Exploring simple siamese representation learning,

    work page Pith review arXiv

  2. [4]

    org/abs/2011.10566

    URLhttps://arxiv. org/abs/2011.10566. Jeongwhan Choi, Hyowon Wi, Jayoung Kim, Yehjin Shin, Kookjin Lee, Nathaniel Trask, and Noseong Park. Graph convolutions enrich the self-attention in transformers!,

    arXiv 2011

  3. [7]

    Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex

    URL https://arxiv.org/abs/2010.11929. Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http:// Skylion007.github.io/OpenWebTextCorpus,

    arXiv 2010

  4. [8]

    Hiroki Nakamura, Masashi Okada, and Tadahiro Taniguchi

    URL https://arxiv.org/abs/1609.02907. Hiroki Nakamura, Masashi Okada, and Tadahiro Taniguchi. Representation uncertainty in self-supervised learning as variational inference. InProceedings of the IEEE/CVF International Conference on Computer Vision, pp. 16484–16493,

    arXiv

  5. [12]

    Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio

    URL https://proceedings.neurips.cc/paper_files/paper/ 2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf. Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks,

    work page 2017

  6. [14]

    Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin D

    URL https://lilianweng.github.io/ posts/2018-06-24-attention/. Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin D. Haeffele, and Yi Ma. White-box transformers via sparse rate reduction,

    work page 2018

  7. [15]

    URL https://arxiv.org/abs/ 2306.01129. 8

    work page Pith review arXiv

  8. [2015]

    doi: 10.1007/s11263-015-0816-y

    ISSN 1573-1405. doi: 10.1007/s11263-015-0816-y. URL https://doi.org/10.1007/s11263-015-0816-y. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.),Advances...

    doi:10.1007/s11263-015-0816-y

Show all 15 references
  1. [2017]

    JacobDevlin, Ming-WeiChang, KentonLee, andKristinaToutanova

    URLhttps://arxiv.org/abs/1707.08819. JacobDevlin, Ming-WeiChang, KentonLee, andKristinaToutanova. Bert: Pre-trainingofdeepbidirectional transformers for language understanding,

    arXiv

  2. [2018]

    Lilian Weng

    URLhttps://arxiv.org/abs/1710.10903. Lilian Weng. Attention? attention! lilianweng.github.io,

    arXiv

  3. [2019]

    URLhttps://arxiv.org/abs/1810.04805. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words...

    arXiv

  4. [2020]

    Hila Chefer, Shir Gur, and Lior Wolf

    URLhttps://arxiv.org/abs/2006.11477. Hila Chefer, Shir Gur, and Lior Wolf. Transformer interpretability beyond attention visualization,

    2006 arXiv

  5. [2021]

    Tianlong Chen, Xiaohan Chen, Wuyang Chen, Howard Heaton, Jialin Liu, Zhangyang Wang, and Wotao Yin

    URL https://arxiv.org/abs/2012.09838. Tianlong Chen, Xiaohan Chen, Wuyang Chen, Howard Heaton, Jialin Liu, Zhangyang Wang, and Wotao Yin. Learning to optimize: A primer and a benchmark,

    2012 arXiv

  6. [2023]

    7 Under review

    URL https://arxiv.org/pdf/2304.07658.pdf. 7 Under review. Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. Imagenet large scale visual recognitio...

    arXiv

  7. [2024]

    Aditya Ravuri, Francisco Vargas, Vidhi Lalchand, and Neil D Lawrence

    URLhttps://arxiv.org/abs/2405.17412. Aditya Ravuri, Francisco Vargas, Vidhi Lalchand, and Neil D Lawrence. Dimensionality reduction as prob- abilistic inference. In Fifth Symposium on Advances in Approximate Bayesian Inference ,

    arXiv

Pith tools

Reviewed August 6, 2026 · model on record in the stance chip above.