Publications

Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs
Tianyu Guo, Druv Pai, Yu Bai, Jiantao Jiao, Michael Jordan, Song Mei
In review at ICLR 2025

Attention-Only Transformers via Unrolled Subspace Denoising
Peng Wang, Yifu Lu, Yaodong Yu, Druv Pai, Qing Qu, Yi Ma
In review at ICLR 2025

Token Statistics Transformer: Linear Time Attention via Variational Rate Reduction
Ziyang Wu, Tianjiao Ding, Druv Pai, Jingyuan Zhang, Weida Wang, Yaodong Yu, Yi Ma, Benjamin Haeffele
In review at ICLR 2025

Scaling White-Box Transformers for Vision
Jinrui Yang, Xianhang Li, Druv Pai, Yuyin Zhou, Yi Ma, Yaodong Yu, Cihang Xie
Accepted at NeurIPS 2024
project website - code

A Geometric Analysis of Maximal Coding Rate Reduction
Peng Wang, Huikang Liu, Druv Pai, Yaodong Yu, Zhihui Zhu, Qing Qu, Yi Ma
Accepted at ICML 2024

White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is?
Yaodong Yu*, Sam Buchanan*, Druv Pai*, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Hao Bai, Yuexiang Zhai, Benjamin Haeffele, Yi Ma
Accepted at JMLR
project website - code

Masked Completion via Structured Diffusion with White-Box Transformers
Druv Pai, Sam Buchanan, Ziyang Wu, Tianzhe Chu, Yaodong Yu, Yi Ma
Accepted at ICLR 2024, accepted at CPAL 2024 (non-archival track)
project website - code

Congestion Pricing for Efficiency and Equity: Theory and Applications to the San Francisco Bay Area
Chinmay Maheshwari, Kshitij Kulkarni, Druv Pai, Jiarui Yang, Manxi Wu, Shankar Sastry

Emergence of Segmentation with Minimalistic White-Box Transformers
Yaodong Yu, Tianzhe Chu, Shengbang Tong, Ziyang Wu, Druv Pai, Sam Buchanan, Yi Ma
Accepted (oral) at CPAL 2024
project website - code

White-Box Transformers via Sparse Rate Reduction
Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin Haeffele, Yi Ma
Accepted (poster) at NeurIPS 2023
project website - code

Representation Learning via Manifold Flattening and Reconstruction
Michael Psenka, Druv Pai, Vishal Raman, Shankar Sastry, Yi Ma
Accepted (poster) at SLowDNN 2023, accepted at JMLR
project website - code

Closed-Loop Transcription via Convolutional Sparse Coding
Xili Dai, Ke Chen, Shengbang Tong, Jingyuan Zhang, Xingjian Gao, Mingyang Li, Druv Pai, Yuexiang Zhai, Xiaojun Yuan, Heung-Yeung Shum, Lionel Ni, Yi Ma
Accepted (poster) at SLowDNN 2023, accepted (oral) at CPAL 2024

Pursuit of a Discriminative Representation for Multiple Subspaces via Sequential Games
Druv Pai, Michael Psenka, Chih-Yuan Chiu, Manxi Wu, Edgar Dobriban, Yi Ma
Accepted (poster) at SLowDNN 2023, accepted at Journal of the Franklin Institute
code

Independent and Decentralized Learning in Markov Potential Games
Chinmay Maheshwari, Manxi Wu, Druv Pai, Shankar Sastry
In review at IEEE Transactions on Automatic Control
code