A Geometric Analysis of Transformer Representations via Optimal Transport
Yadnyesh Chakane, Sunishka Sharma, Vishnu Vardhan Lanka, Janhavi Khindkar · Team Optimal Transport
Submitted to AI Safety x Physics Grand Challenge. Sprint projects are early-stage work by participants, not Apart Research publications.
Understanding the internal workings of transformers is a major challenge in deep learning; while these models achieve state-of-the-art performance, their multi-layer architectures operate as "black boxes," obscuring the principles that guide their information processing pathways. To address this, we used Optimal Transport (OT) to analyze the geometric transformations of representations between layers in both trained and untrained transformer models. By treating layer activations as empirical distributions, we computed layer-to-layer OT distances to quantify the extent of geometric rearrangement, complementing this with representation entropy measurements to track information content. Our results show that trained models exhibit a structured, three-phase information processing strategy (encode-refine-decode), characterized by an information bottleneck. This is evident from a U-shaped OT distance profile, where high initial costs give way to a low-cost "refinement" phase before a final, high-cost projection to the output layer. This structure is entirely absent in untrained models, which instead show chaotic, uniformly high-cost transformations. We conclude that OT provides a powerful tool for revealing the learned, efficient information pathways in neural networks, demonstrating that learning is not merely about fitting data, but about creating an organized, information-theoretically efficient pipeline to process representations.
Reviews
Cool! Interesting and sensible idea (I enjoyed reading Tishby's information bottleneck treatments of NNs back in the day, but am not aware of anyone trying to use OT to study information flow within a model). Solid execution for a quick hackathon project. Also very clearly written, I appreciated how easy it was to read.
My main concern is just that it's underpowered, and the results are still very sparse.. I think it'd be very valuable to look at how this changes over the course of training and to look at some smaller models where we have some ground-truth.
Love this project and the simple conclusion of an encoding and a refinement phase. I would imagine this is a generalizable statement across NNs (worth a test) and it's interesting that the minimization of work isn't a *necessary* condition for information processing but that the refinement phase is an iterative improvement on the information as it passes through the network to reduce how much noise is output in latter layers. It generally looks to be a linear-ish increase in work between layers as it improves the compressed representation of the concepts it runs over, though it would be very interesting to see whether we can use this to study suddenly-changed models and see a non-linear change in work over the trained networks, especially over their checkpoints (models saved at intermediate intervals of the training process). Great work, very simple, good formalization, and a good introduction. Extensions on this work will show whether it'll be useful for AI safety but it has some interesting implications for how models work to create their output in general.
Read full reviewShow less
This project uses methods based on optimal transport to analyze the layer-wise evolution of representations in transformers. The project is interesting and well-executed, and the main ideas are explained well. The project could be improved by more explicitly articulating the connection to AI safety.
Cite this project
@misc{chakane2025geometric,
title = {{A Geometric Analysis of Transformer Representations via Optimal Transport}},
author = {Yadnyesh Chakane and Sunishka Sharma and Vishnu Vardhan Lanka and Janhavi Khindkar},
year = {2025},
month = jul,
note = {Submitted to AI Safety x Physics Grand Challenge, an Apart Research Sprint},
howpublished = {\url{https://apartresearch.com/sprints/projects/a-geometric-analysis-of-transformer-representations-via-optimal-transport-qjdf}},
url = {https://apartresearch.com/sprints/projects/a-geometric-analysis-of-transformer-representations-via-optimal-transport-qjdf}
}More from AI Safety x Physics Grand Challenge
- Prize winnerView project: Constrained Belief Updates and Geometric Structures in Transformer Representations for the RRXOR Process
Constrained Belief Updates and Geometric Structures in Transformer Representations for the RRXOR Process
In this work, we seek to further understand the computational structure formed by transformers trained to predict the next token of a data set. In particular, we extend the work of (Piotrowski et al., 2025) by analyzing …
- Prize winnerView project: Momentum–Point-Perplexity Mechanics in Large Language Models=
Momentum–Point-Perplexity Mechanics in Large Language Models=
Fractal
This work analyzes the hidden states of twenty different open-source transformer language models, ranging from small to medium size and covering five major architectures. The key discovery is that these models show …
- Prize winnerView project: AI agentic system epidemiology
AI agentic system epidemiology
XIA-for-Med
As AI systems scale into decentralized, multi-agent deployments, emergent vulnerabilities challenge our ability to evaluate and manage systemic risks. In this work, we adapt classical epidemiological modeling …