← Search

Song Mei

27 accepted papers

2025

Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformers

NeurIPS 2025poster

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 also prone to hallucinating incorrect information. However, the reasons for this phenomenon remain poorly understood. In…

Cited by 0SourceScholar
2025

Implicit Bias of Gradient Descent for Non-Homogeneous Deep Networks

ICML 2025poster

We establish the asymptotic implicit bias of gradient descent (GD) for generic non-homogeneous deep networks under exponential loss. Specifically, we characterize three key properties of GD iterates starting from a sufficiently small empirical risk, where the threshold is determined by a measure of…

Cited by 0SourcePDFScholar
2025

Improving LLM Safety Alignment with Dual-Objective Optimization

ICML 2025poster

Existing training-time safety alignment techniques for large language models (LLMs) remain vulnerable to jailbreak attacks. Direct preference optimization (DPO), a widely deployed alignment method, exhibits limitations in both experimental and theoretical contexts as its loss function proves subopti…

2025

OVERT: A Benchmark for Over-Refusal Evaluation on Text-to-Image Models

NeurIPS 2025poster

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 prevent harmful outputs, they often lead to overly cautious behavior ---rejecting even benign prompts---a phenomenon known as…

Cited by 0SourcecodeScholar
2025

Simplicity Prevails: Rethinking Negative Preference Optimization for LLM Unlearning

NeurIPS 2025poster

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 utility. Despite the increasing demand for unlearning, a technically-grounded optimization framework is lacking. Gradient as…

Cited by 0SourcecodeScholar
2025

U-Nets as Belief Propagation: Efficient Classification, Denoising, and Diffusion in Generative Hierarchical Models

ICLR 2025poster

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 diffusion modeling. However, a theoretical explanation of the U-Net architecture design has not yet been fully established.…

Cited by 10SourcePDFScholar
2024

How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations

ICLR 2024poster

While large language models based on the transformer architecture have demonstrated remarkable in-context learning (ICL) capabilities, understandings of such capabilities are still in an early stage, where existing theory and mechanistic understanding focus mostly on simple scenarios such as learnin…

Cited by 61SourcePDFScholar
2024

Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization

NeurIPS 2024poster

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 in the first phase but decreases monotonically in the second phase. We investigate this phenomenon in two-layer networks…

Cited by 6SourcePDFScholar
2024

Statistical Estimation in the Spiked Tensor Model via the Quantum Approximate Optimization Algorithm

NeurIPS 2024spotlight

The quantum approximate optimization algorithm (QAOA) is a general-purpose algorithm for combinatorial optimization that has been a promising avenue for near-term quantum advantage. In this paper, we analyze the performance of the QAOA on the spiked tensor model, a statistical estimation problem th…

Cited by 1SourcePDFScholar
2024

Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining

ICLR 2024poster

Large transformer models pretrained on offline reinforcement learning datasets have demonstrated remarkable in-context reinforcement learning (ICRL) capabilities, where they can make good decisions when prompted with interaction trajectories from unseen environments. However, when and how transforme…

2023

Partially Observable RL with B-Stability: Unified Structural Condition and Sharp Sample-Efficient Algorithms

ICLR 2023top-25%

Partial Observability---where agents can only observe partial information about the true underlying state of the system---is ubiquitous in real-world applications of Reinforcement Learning (RL). Theoretically, learning a near-optimal policy under partial observability is known to be hard in the wors…

Cited by 31SourcePDFScholar
2023

Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection

NeurIPS 2023oral

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 with training and test examples, without any parameter update to the model. This work first provides a comprehensive st…

2023

What can a Single Attention Layer Learn? A Study Through the Random Features Lens

NeurIPS 2023poster

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 in modern artificial intelligence. This paper presents a rigorous theoretical study on the learning and generalization of…

Cited by 34SourcePDFScholar
2022

Efficient Phi-Regret Minimization in Extensive-Form Games via Online Mirror Descent

NeurIPS 2022accept

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-the-art techniques and analyses in NFGs to learning EFGs, but typically suffers from computational intractability due to…

Cited by 25SourcePDFScholar
2022

Efficient and Differentiable Conformal Prediction with General Function Classes

ICLR 2022poster

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 for learned prediction sets are \emph{valid coverage} and \emph{good efficiency} (such as low length or low cardinality).…

2022

Near-Optimal Learning of Extensive-Form Games with Imperfect Information

ICML 2022spotlight

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 of algorithms that require only $\widetilde{\mathcal{O}}((XA+YB)/\varepsilon^2)$ episodes of play to find an $\varepsilon…

Cited by 36SourcePDFScholar
2022

When Can We Learn General-Sum Markov Games with a Large Number of Players Sample-Efficiently?

ICLR 2022poster

Multi-agent reinforcement learning has made substantial empirical progresses in solving games with a large number of players. However, theoretically, the best known sample complexity for finding a Nash equilibrium in general-sum games scales exponentially in the number of players due to the size of…

Cited by 124SourcePDFScholar
2021

Don’t Just Blame Over-parametrization for Over-confidence: Theoretical Analysis of Calibration in Binary Classification

ICML 2021spotlight

Modern machine learning models with high accuracy are often miscalibrated—the predicted top probability does not reflect the actual accuracy, and tends to be \emph{over-confident}. It is commonly believed that such over-confidence is mainly due to \emph{over-parametrization}, in particular when the…

Cited by 64SourcePDFScholar
2021

Exact Gap between Generalization Error and Uniform Convergence in Random Feature Models

ICML 2021spotlight

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 as deep neural networks. To better understand this gap, we study the uniform convergence in the nonlinear random feature…

Cited by 27SourcePDFScholar
2021

Understanding the Under-Coverage Bias in Uncertainty Estimation

NeurIPS 2021spotlight

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 frequently observed that quantile regression---a vanilla algorithm for learning quantiles with asymptotic guarantees---tends to…

Cited by 18SourcePDFScholar
2020

When Do Neural Networks Outperform Kernel Methods?

NeurIPS 2020poster

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 space (RKHS) methods. Recent empirical work showed that, for some classification tasks, RKHS methods can replace NNs with…