@article{zhang2026mpt,title={A Theoretical Framework for Masked Pretraining (MPT)},author={Zhang, Qi and Zhou, Runyu and {<u>Yifei Wang</u>} and Wang#, Yisen},journal={Journal of Machine Learning Research},year={2026}}
NeurIPS
Quest: Training Frontier Deep Research Agents with Fully Synthetic Tasks
Jian Xie, Tianhe Lin, Zilu Wang, Yuting Ning, Yuekun Yao, Tianci Xue, Zhehao Zhang, Zhongyang Li, Kai Zhang, Yufan Wu, Shijie Chen, Boyu Gou, Mingzhe Han, Yifei Wang, Vint Lee, Xinpeng Wei, Xiangjun Wang, Yu Su, and Huan Sun
Lizhe Fang*, Yifei Wang*, Khashayar Gatmiry, Lei Fang, and Yisen Wang
ICLR, 2025
We discovered an expressive invariant in-context learning scheme (InvICL) that achieves permutation invariance of in-context demonstrations while preserving autoregressive nature and full context awareness at the same time.
Can In-context Learning Really Generalize to Out-of-distribution Tasks?
Qixun Wang, Yifei Wang, Yisen Wang, and Xianghua Ying
ICLR, 2025
With controlled experiments, we found that in-context learning still happens only with in-domain tasks and hardly generalizes to novel OOD tasks. In other words, LLMs’ in-context abilities are learned essentially through training data with likewise tasks.
Beyond Interpretability: The Gains of Feature Monosemanticity on Model Robustness
Qi Zhang*, Yifei Wang*, Jingyi Cui, Xiang Pan, Qi Lei, Stefanie Jegelka, and Yisen Wang
ICLR, 2025
We found that the merits of feature monosemanticity (as studied in mechanistic interpretability) extend beyond interpretability to improving robustness across various challenges like noisy data, limited training examples, and such.
Projection Head is Secretly an Information Bottleneck
Zhuo Ouyang, Kaiwen Hu, Qi Zhang, Yifei Wang, and Yisen Wang
ICLR, 2025
We showed that projection heads serve as an information bottleneck that prevent features from collapsing toward the pretraining task (e.g. instance classification).
A Theoretical Understanding of Self-Correction through In-context Alignment
Yifei Wang*, Yuyang Wu*, Zeming Wei, Stefanie Jegelka, and Yisen Wang
In NeurIPS, 2024
🏆 Best Paper Award at ICML 2024 ICL Workshop We proposed the first theoretical explanation of how LLM self-correction works (as in OpenAI o1) and showed its effectiveness against social bias and jailbreak attacks.
Understanding the Role of Equivariance in Self-supervised Learning
Yifei Wang*, Kaiwen Hu*, Sharut Gupta, Ziyu Ye, Yisen Wang, and Stefanie Jegelka
In NeurIPS, 2024
First theoretical explanation of how equivariant prediction helps self-supervised representation learning – using information theory & probabilistic graphical models.
In-Context Symmetries: Self-Supervised Learning through Contextual World Models
Sharut Gupta*, Chenyu Wang*, Yifei Wang*, Tommi Jaakkola, and Stefanie Jegelka
In NeurIPS, 2024
Oral Presentation (top 4) at NeurIPS 2024 SSL Workshop & featured by MIT CSAIL News 📰 We introduced unsupervised test-time adaptation ability to self-supervised learning through a contextual world model designed for joint embedding (JEPA) models.
Reasoning in Reasoning: A Hierarchical Framework for Better and Faster Neural Theorem Proving
Ziyu Ye, Jiacheng Chen, Jonathan Light, Yifei Wang, Jiankai Sun, Mac Schwager, Philip Torr, Guohao Li, Yuxin Chen, Kaiyu Yang, Yisong Yue, and Ziniu Hu
In NeurIPS 2024 Workshop on Mathematical Reasoning and AI, 2024
Inspired by NMF, we introduced a simple technique (one-line) that attains 90% feature sparsity and 10x feature interpretability for self-supervised contrastive learning, with theoretical guarantees on its disentanglement and performance.
Do Generated Data Always Help Contrastive Learning?
Yifei Wang*, Jizhe Zhang*, and Yisen Wang
In ICLR, 2024
We revealed both theoretically and practically that synthetic data introduces fundamental bias to SSL generalization, but, with an adaptive strategy of data mixing and augmentation, can yield substantial benefits.
On the Generalization of Multi-modal Contrastive Learning
Qi Zhang*, Yifei Wang*, and Yisen Wang
In ICML, 2023
We established the first generalization analysis for multi-modal contrastive learning (e.g., CLIP) and explained how it outperforms self-supervised contrastive learning.
A Message Passing Perspective on Learning Dynamics of Contrastive Learning
Yifei Wang*, Qi Zhang*, Tianqi Du, Jiansheng Yang, Zhouchen Lin, and Yisen Wang
In ICLR, 2023
We revealed that contrastive learning performs message passing on sample graph, which connects self-supervised learning and graph neural networks as a whole.
Towards a Unified Theoretical Understanding of Non-contrastive Learning via Rank Differential Mechanism
Zhijian Zhuo*, Yifei Wang*, Jinwen Ma, and Yisen Wang
In ICLR, 2023
We revealed that various asymmtric designs in non-contrastive learning (BYOL, SimSiam, DINO, SwAV) can be explained from a unified spectral filtering perspective.
Rethinking the Effect of Data Augmentation in Adversarial Contrastive Learning
Rundong Luo*, Yifei Wang*, and Yisen Wang
In ICLR, 2023
We improved adversarial robustness under AutoAttack by 9% in the unsupervised setting with a dynamic training schedule, without extra computation cost.
How Mask Matters: Towards Theoretical Understandings of Masked Autoencoders
Qi Zhang*, Yifei Wang*, and Yisen Wang
In NeurIPS Spotlight (Top 5%), 2022
We theoretically explained how masked autoencoders work and revealed their mathematical connections to joint embedding methods, unifying them as a whole. Extended version published in JMLR 2026.
Chaos is a Ladder: A New Theoretical Understanding of Contrastive Learning via Augmentation Overlap
Yifei Wang*, Qi Zhang*, Yisen Wang, Jiansheng Yang, and Zhouchen Lin
In ICLR, 2022
A new augmentation overlap theory for understanding the generalization of contrastive learning. Cited over 150 times. Extended version was accepted at JMLR.
A Unified Contrastive Energy-based Model for Understanding the Generative Ability of Adversarial Training
Yifei Wang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin
In ICLR, 2022
🏆 Silver Best Paper Award at ICML 2021 AdvML workshop From an energy-based perspective, we formulated contrastive learning as a generative model, and established the connection between adversarial training and maximum likelihood, thus briding generative and discriminative models together.
Reparameterized Sampling for Generative Adversarial Networks
Yifei Wang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin
In ECML-PKDD, 2021
🏆 Best ML Paper Award (1/685), invited to Machine Learning We explored using GAN discriminator (as a good reward model) to bootstrap sample quality through an efficient MCMC algorithm, which not only guarantees theoretical convergence but also improves sample efficiency and quality in practice.