Research
I work on the mathematical theory of transformers — modeling tokens as interacting particles to understand how attention moves, clusters, and represents information. I work with Yury Polyanskiy and Philippe Rigollet at MIT.
Papers
Continuous First, Discrete Later: VQ-VAEs Without Dimensional Collapse
Trained VQ-VAE representations tend to collapse into a tiny subspace — 1–2% of full rank — and we show this dimensional collapse imposes a hard lower bound on the loss that codebook-improvement tricks cannot overcome. The fix is simple: warm up the model as an ordinary (unquantized) autoencoder before switching on vector quantization. On VQGAN and WavTokenizer, this warm-up restores representation dimension and improves reconstruction and perceptual quality at the same training budget, and our theory predicts the right moment to switch.
BibTeX
@misc{zhao2026continuousfirstdiscretelater,
title={Continuous First, Discrete Later: VQ-VAEs Without Dimensional Collapse},
author={Xinyu Zhao and Nikita Karagodin and Hamed Hassani and Sinan Hersek and Paul Pu Liang and Yury Polyanskiy},
year={2026},
eprint={2605.06870},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2605.06870},
}Normalization in Attention Dynamics
We model token representations in a deep transformer as interacting particles on the sphere and show that normalization acts as a form of speed regulation on their dynamics. This perspective gives a unified analysis of Post-LN, Pre-LN, Mix-LN, Peri-LN, and nGPT, explaining how each scheme shapes clustering and representation collapse across layers. The comparison identifies Peri-LN as a particularly effective choice.
BibTeX
@misc{karagodin2025normalizationattentiondynamics,
title={Normalization in Attention Dynamics},
author={Nikita Karagodin and Shu Ge and Yury Polyanskiy and Philippe Rigollet},
year={2025},
eprint={2510.22026},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2510.22026},
}Clustering in Causal Attention Masking
This is the first rigorous treatment of causally masked attention dynamics — the version transformers actually use for generation, which breaks the mean-field gradient-flow structure earlier analyses relied on. We prove that tokens converge to a single cluster for arbitrary key-query matrices with identity value matrix, going well beyond prior results that required all three matrices to be scaled identities. A connection to the classical Rényi parking problem takes first theoretical steps toward explaining the meta-stable states observed in practice.
BibTeX
@misc{karagodin2024clusteringcausalattentionmasking,
title={Clustering in Causal Attention Masking},
author={Nikita Karagodin and Yury Polyanskiy and Philippe Rigollet},
year={2024},
eprint={2411.04990},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2411.04990},
}Other fascinating things
Papers and ideas I didn't write but find remarkable.
The Unreasonable Effectiveness of Mathematics
The classic essay on why mathematical concepts developed in pure abstraction turn out to describe physical reality with uncanny precision.