Publications (61)
Proximal algorithms for constrained composite optimization, with applications to solving low-rank SDPs
Yu Bai, John Duchi, Song Mei
We study a family of (potentially non-convex) constrained optimization problems with convex composite structure. Through a novel analysis of non-smooth geometry, we show that proxi…
When Can We Learn General-Sum Markov Games with a Large Number of Players Sample-Efficiently?
Ziang Song, Song Mei, Yu Bai
Multi-agent reinforcement learning has made substantial empirical progresses in solving games with a large number of players. However, theoretically, the best known sample complexi…
Efficient and Differentiable Conformal Prediction with General Function Classes
Yu Bai, Song Mei, Huan Wang +2
Quantifying the data uncertainty in learning tasks is often done by learning a prediction interval or prediction set of the label given the input. Two commonly desired properties f…
Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs
Tianyu Guo, Druv Pai, Yu Bai +3
Practitioners have consistently observed three puzzling phenomena in transformer-based large language models (LLMs): attention sinks, value-state drains, and residual-state peaks,…
Uncertainty Intervals for Prediction Errors in Time Series Forecasting
Hui Xu, Song Mei, Stephen Bates +2
Inference for prediction errors is critical in time series forecasting pipelines. However, providing statistically meaningful uncertainty intervals for prediction errors remains re…
Improving LLM Safety Alignment with Dual-Objective Optimization
Xuandong Zhao, Will Cai, Tianneng Shi +4
Existing training-time safety alignment techniques for large language models (LLMs) remain vulnerable to jailbreak attacks. Direct preference optimization (DPO), a widely deployed…
Negative Preference Optimization: From Catastrophic Collapse to Effective Unlearning
Ruiqi Zhang, Licong Lin, Yu Bai +1
Large Language Models (LLMs) often memorize sensitive, private, or copyrighted data during pre-training. LLM unlearning aims to eliminate the influence of undesirable data from the…
An Overview of Diffusion Models: Applications, Guided Generation, Statistical Rates and Optimization
Minshuo Chen, Song Mei, Jianqing Fan +1
Diffusion models, a powerful and universal generative AI technology, have achieved tremendous success in computer vision, audio, reinforcement learning, and computational biology.…
How Do LLMs Perform Two-Hop Reasoning in Context?
Tianyu Guo, Hanlin Zhu, Ruiqi Zhang +4
``Socrates is human. All humans are mortal. Therefore, Socrates is mortal.'' This form of argument illustrates a typical pattern of two-hop reasoning. Formally, two-hop reasoning r…
GSM-Agent: Understanding Agentic Reasoning Using Controllable Environments
Hanlin Zhu, Tianyu Guo, Song Mei +4
As LLMs are increasingly deployed as agents, agentic reasoning - the ability to combine tool use, especially search, and reasoning - becomes a critical skill. However, it is hard t…
The Landscape of Empirical Risk for Non-convex Losses
Song Mei, Yu Bai, Andrea Montanari
Most high-dimensional estimation and prediction methods propose to minimize a cost function (empirical risk) that is written as a sum of losses associated to each data point. In th…
Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining
Licong Lin, Yu Bai, Song Mei
Large transformer models pretrained on offline reinforcement learning datasets have demonstrated remarkable in-context reinforcement learning (ICRL) capabilities, where they can ma…
Generalization error of random features and kernel methods: hypercontractivity and kernel matrix concentration
Song Mei, Theodor Misiakiewicz, Andrea Montanari
Consider the classical supervised learning problem: we are given data , , with a response and a covariate…
U-Nets as Belief Propagation: Efficient Classification, Denoising, and Diffusion in Generative Hierarchical Models
Song Mei
U-Nets are among the most widely used architectures in computer vision, renowned for their exceptional performance in applications such as image segmentation, denoising, and diffus…
SPEED-RL: Faster Training of Reasoning Models via Online Curriculum Learning
Ruiqi Zhang, Daman Arora, Song Mei +1
Training large language models with reinforcement learning (RL) against verifiable rewards significantly enhances their reasoning abilities, yet remains computationally expensive d…
Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization
Yuhang Cai, Jingfeng Wu, Song Mei +2
The typical training of neural networks using large stepsize gradient descent (GD) under the logistic loss often involves two distinct phases, where the empirical risk oscillates i…
Lower Bounds for Learning in Revealing POMDPs
Fan Chen, Huan Wang, Caiming Xiong +2
This paper studies the fundamental limits of reinforcement learning (RL) in the challenging \emph{partially observable} setting. While it is well-established that learning in Parti…
The landscape of the spiked tensor model
Gerard Ben Arous, Song Mei, Andrea Montanari +1
We consider the problem of estimating a large rank-one tensor , in Gaussian noise. Earlier work characterized…
A Statistical Theory of Contrastive Pre-training and Multimodal Generative AI
Kazusato Oko, Licong Lin, Yuhang Cai +1
Multi-modal generative AI systems, such as those combining vision and language, rely on contrastive pre-training to learn representations across different modalities. While their p…
Near-Optimal Learning of Extensive-Form Games with Imperfect Information
Yu Bai, Chi Jin, Song Mei +1
This paper resolves the open question of designing near-optimal algorithms for learning imperfect-information extensive-form games from bandit feedback. We present the first line o…
A Mean Field View of the Landscape of Two-Layers Neural Networks
Song Mei, Andrea Montanari, Phan-Minh Nguyen
Multi-layer neural networks are among the most powerful models in machine learning, yet the fundamental reasons for this success defy mathematical understanding. Learning a neural…
DiscoLoop: Looping Discrete Embeddings and Continuous Hidden States for Multi-hop Reasoning
Hengyu Fu, Tianyu Guo, Zixuan Wang +5
Large language models achieve strong performance on many reasoning tasks when allowed to externalize intermediate steps as Chain-of-Thought (CoT). However, many questions require t…
Local convexity of the TAP free energy and AMP convergence for Z2-synchronization
Michael Celentano, Zhou Fan, Song Mei
We study mean-field variational Bayesian inference using the TAP approach, for Z2-synchronization as a prototypical example of a high-dimensional Bayesian model. We show that for a…
Learning with invariances in random features and kernel models
Song Mei, Theodor Misiakiewicz, Andrea Montanari
A number of machine learning tasks entail a high degree of invariance: the data distribution does not change if we act on the data with a certain group of transformations. For inst…
Performance and limitations of the QAOA at constant levels on large sparse hypergraphs and spin glass models
Joao Basso, David Gamarnik, Song Mei +1
The Quantum Approximate Optimization Algorithm (QAOA) is a general purpose quantum algorithm designed for combinatorial optimization. We analyze its expected performance and prove…
Efficient Phi-Regret Minimization in Extensive-Form Games via Online Mirror Descent
Yu Bai, Chi Jin, Song Mei +2
A conceptually appealing approach for learning Extensive-Form Games (EFGs) is to convert them to Normal-Form Games (NFGs). This approach enables us to directly translate state-of-t…
An Overview of Large Language Models for Statisticians
Wenlong Ji, Weizhe Yuan, Emily Getzen +7
Large Language Models (LLMs) have emerged as transformative tools in artificial intelligence (AI), exhibiting remarkable capabilities across diverse tasks such as text generation,…
Don't Just Blame Over-parametrization for Over-confidence: Theoretical Analysis of Calibration in Binary Classification
Yu Bai, Song Mei, Huan Wang +1
Modern machine learning models with high accuracy are often miscalibrated -- the predicted top probability does not reflect the actual accuracy, and tends to be over-confident. It…
Implicit Bias of Gradient Descent for Non-Homogeneous Deep Networks
Yuhang Cai, Kangjie Zhou, Jingfeng Wu +3
We establish the asymptotic implicit bias of gradient descent (GD) for generic non-homogeneous deep networks under exponential loss. Specifically, we characterize three key propert…
Deep Networks as Denoising Algorithms: Sample-Efficient Learning of Diffusion Models in High-Dimensional Graphical Models
Song Mei, Yuchen Wu
We investigate the approximation efficiency of score functions by deep neural networks in diffusion-based generative modeling. While existing approximation theories utilize the smo…
On a molecular based Q-tensor model for liquid crystals with density variations
Song Mei, Pingwen Zhang
In this article, we study the new Q-tensor model previously derived from Onsager's molecular theory by Han \textit{et al.} [Arch. Rational Mech. Anal., 215.3 (2014), pp. 741-809] f…
Statistical Estimation in the Spiked Tensor Model via the Quantum Approximate Optimization Algorithm
Leo Zhou, Joao Basso, Song Mei
The quantum approximate optimization algorithm (QAOA) is a general-purpose algorithm for combinatorial optimization. In this paper, we analyze the performance of the QAOA on a stat…
Partially Observable RL with B-Stability: Unified Structural Condition and Sharp Sample-Efficient Algorithms
Fan Chen, Yu Bai, Song Mei
Partial Observability -- where agents can only observe partial information about the true underlying state of the system -- is ubiquitous in real-world applications of Reinforcemen…
The generalization error of random features regression: Precise asymptotics and double descent curve
Song Mei, Andrea Montanari
Deep learning methods operate in regimes that defy the traditional statistical mindset. Neural network architectures often contain more parameters than training samples, and are so…
A Statistical Theory of Contrastive Learning via Approximate Sufficient Statistics
Licong Lin, Song Mei
Contrastive learning -- a modern approach to extract useful representations from unlabeled data by training models to distinguish similar samples from dissimilar ones -- has driven…
Simplicity Prevails: Rethinking Negative Preference Optimization for LLM Unlearning
Chongyu Fan, Jiancheng Liu, Licong Lin +4
This work studies the problem of large language model (LLM) unlearning, aiming to remove unwanted data influences (e.g., copyrighted or harmful content) while preserving model util…
OVERT: A Benchmark for Over-Refusal Evaluation on Text-to-Image Models
Ziheng Cheng, Yixiao Huang, Hui Xu +4
Text-to-Image (T2I) models have achieved remarkable success in generating visual content from text inputs. Although multiple safety alignment strategies have been proposed to preve…
Exact Gap between Generalization Error and Uniform Convergence in Random Feature Models
Zitong Yang, Yu Bai, Song Mei
Recent work showed that there could be a large gap between the classical uniform convergence bound and the actual test error of zero-training-error predictors (interpolators) such…
Mean-field variational inference with the TAP free energy: Geometric and statistical properties in linear models
Michael Celentano, Zhou Fan, Licong Lin +1
We study mean-field variational inference in a Bayesian linear model when the sample size n is comparable to the dimension p. In high dimensions, the common approach of minimizing…
Sample-Efficient Learning of Correlated Equilibria in Extensive-Form Games
Ziang Song, Song Mei, Yu Bai
Imperfect-Information Extensive-Form Games (IIEFGs) is a prevalent model for real-world games involving imperfect information and sequential plays. The Extensive-Form Correlated Eq…
OpenAI GPT-5 System Card
Aaditya Singh, Adam Fry, Adam Perelman +483
This is the system card published alongside the OpenAI GPT-5 launch, August 2025. GPT-5 is a unified system with a smart and fast model that answers most questions, a deeper reason…
Near-optimal multiple testing in Bayesian linear models with finite-sample FDR control
Taejoo Ahn, Licong Lin, Song Mei
In high dimensional variable selection problems, statisticians often seek to design multiple testing procedures that control the False Discovery Rate (FDR), while concurrently iden…
When Do Neural Networks Outperform Kernel Methods?
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz +1
For a certain scaling of the initialization of stochastic gradient descent (SGD), wide neural networks (NN) have been shown to be well approximated by reproducing kernel Hilbert sp…
Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformers
Yixiao Huang, Hanlin Zhu, Tianyu Guo +5
Large language models (LLMs) can acquire new knowledge through fine-tuning, but this process exhibits a puzzling duality: models can generalize remarkably from new facts, yet are a…
Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection
Yu Bai, Fan Chen, Huan Wang +2
Neural sequence models based on the transformer architecture have demonstrated remarkable \emph{in-context learning} (ICL) abilities, where they can perform new tasks when prompted…
How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations
Tianyu Guo, Wei Hu, Song Mei +4
While large language models based on the transformer architecture have demonstrated remarkable in-context learning (ICL) capabilities, understandings of such capabilities are still…
Understanding the Under-Coverage Bias in Uncertainty Estimation
Yu Bai, Song Mei, Huan Wang +1
Estimating the data uncertainty in regression tasks is often done by learning a quantile function or a prediction interval of the true label conditioned on the input. It is frequen…
Linearized two-layers neural networks in high dimension
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz +1
We consider the problem of learning an unknown function on the -dimensional sphere with respect to the square loss, given i.i.d. samples $\{(y_i,{\boldsymbol x}_i)\}…
Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit
Song Mei, Theodor Misiakiewicz, Andrea Montanari
We consider learning two layer neural networks using stochastic gradient descent. The mean-field description of this learning dynamics approximates the evolution of the network wei…
Solving SDPs for synchronization and MaxCut problems via the Grothendieck inequality
Song Mei, Theodor Misiakiewicz, Andrea Montanari +1
A number of statistical estimation problems can be addressed by semidefinite programs (SDP). While SDPs are solvable in polynomial time using interior point methods, in practice ge…
Unified Algorithms for RL with Decision-Estimation Coefficients: PAC, Reward-Free, Preference-Based Learning, and Beyond
Fan Chen, Song Mei, Yu Bai
Modern Reinforcement Learning (RL) is more than just learning the optimal policy; Alternative learning goals such as exploring the environment, estimating the underlying model, and…
A Novel Location Free Link Prediction in Multiplex Social Networks
Song Mei, Cong Zhen
In recent decades, the emergence of social networks has enabled internet service providers (e.g., Facebook, Twitter and Uber) to achieve great commercial success. Link prediction i…
Plausible deniability for privacy-preserving data synthesis
Song Mei, Zhiqiang Ye
In the field of privacy protection, publishing complete data (especially high-dimensional data sets) is one of the most challenging problems. The common encryption technology can n…
The Three Stages of Learning Dynamics in High-Dimensional Kernel Methods
Nikhil Ghosh, Song Mei, Bin Yu
To understand how deep learning works, it is crucial to understand the training dynamics of neural networks. Several interesting hypotheses about these dynamics have been made base…
TAP free energy, spin glasses, and variational inference
Zhou Fan, Song Mei, Andrea Montanari
We consider the Sherrington-Kirkpatrick model of spin glasses with ferromagnetically biased couplings. For a specific choice of the couplings mean, the resulting Gibbs measure is e…
Transformers Provably Learn to Internalize Chain-of-Thought
Yixiao Huang, Hanlin Zhu, Zixuan Wang +4
Chain-of-Thought (CoT) prompting substantially improves the sample efficiency of transformers, reducing the complexity of tasks like parity learning from exponential to polynomial…
What can a Single Attention Layer Learn? A Study Through the Random Features Lens
Hengyu Fu, Tianyu Guo, Yu Bai +1
Attention layers -- which map a sequence of inputs to a sequence of outputs -- are core building blocks of the Transformer architecture which has achieved significant breakthroughs…
Learning with convolution and pooling operations in kernel methods
Theodor Misiakiewicz, Song Mei
Recent empirical work has shown that hierarchical convolutional kernels inspired by convolutional neural networks (CNNs) significantly improve the performance of kernel methods in…
Limitations of Lazy Training of Two-layers Neural Networks
Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz +1
We study the supervised learning problem under either of the following two models: (1) Feature vectors are -dimensional Gaussians and responses are $y_i = f_…
gpt-oss-120b & gpt-oss-20b Model Card
OpenAI, :, Sandhini Agarwal +124
We present gpt-oss-120b and gpt-oss-20b, two open-weight reasoning models that push the frontier of accuracy and inference cost. The models use an efficient mixture-of-expert trans…
Analysis of Sequential Quadratic Programming through the Lens of Riemannian Optimization
Yu Bai, Song Mei
We prove that a "first-order" Sequential Quadratic Programming (SQP) algorithm for equality constrained optimization has local linear convergence with rate , where $Î…