Reproduce “Transformers as Unrolled Inference in Probabilistic Laplacian Eigenmaps: An Interpretation and Potential Improvements” on OpenResearch

About this paper

Aditya Ravuri, Neil D. Lawrence

Transformer blocks are interpreted as unrolled inference steps within a probabilistic Laplacian Eigenmaps model, connecting their operations to dimensionality reduction. This theoretical grounding motivates an architectural modification where the attention matrix is replaced by a graph diffusion step (A - I), which improves validation performance and accelerates convergence across language and vision tasks.

Key insights

• Transformer blocks, especially at initialization, perform linear dimensionality reduction by unrolling gradient descent on a probabilistic Laplacian Eigenmaps objective.
• The attention matrix in transformers, when derived from a probabilistic framework, naturally yields a graph Laplacian term (A - I), suggesting a graph diffusion or smoothing operation.
• The linear weight matrices within transformer blocks can be interpreted as learning rotation and scaling factors, corresponding to adaptive step sizes in the unrolled optimization process.

Summary by alphaXiv.