See full list on Google Scholar
2025
-
Beyond Matryoshka: Revisiting Sparse Coding for Adaptive Representation
Tiansheng Wen*, Yifei Wang*, Zequn Zeng, Zhong Peng, Yudi Su, Xinyang Liu, Bo Chen, Hongwei Liu, Stefanie Jegelka, and Chenyu You
ICML Oral Presentation (1%), 2025
-
On the Emergence of Position Bias in Transformers
Xinyi Wu, Yifei Wang, Stefanie Jegelka, and Ali Jadbabaie
ICML, 2025
-
Long-Short Alignment for Effective Long-Context Modeling in LLMs
Tianqi Du*, Haotian Huang*, Yifei Wang, and Yisen Wang
ICML, 2025
-
G1: Teaching LLMs to Reason on Graphs with Reinforcement Learning
Xiaojun Guo*, Ang Li*, Yifei Wang*, Stefanie Jegelka, and Yisen Wang
arXiv preprint arXiv:2505.18499, 2025
-
ICLR Workshop Best Paper Runner-up When More is Less: Understanding Chain-of-Thought Length in LLMs
Yuyang Wu*, Yifei Wang*, Ziyu Ye, Tianqi Du, Stefanie Jegelka, and Yisen Wang
ICLR 2025 Workshop on Reasoning and Planning for LLMs, 2025
🏆 Best Paper Runner-up Award
-
ICLR LLM training and eval What is Wrong with Perplexity for Long-context Language Modeling?
Lizhe Fang*, Yifei Wang*, Zhaoyang Liu, Chenheng Zhang, Stefanie Jegelka, Jinyang Gao, Bolin Ding, and Yisen Wang
ICLR, 2025
-
Rethinking Invariance in In-context Learning
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.
-
2024
-
NeurIPS Best Paper Award
at ICML-W’24 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.
-
NeurIPS Oral at NeurIPS-W’24 Featured by MIT 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 📰 We introduced unsupervised test-time adaptation ability to self-supervised learning through a contextual world model designed for joint embedding (JEPA) models.
-
A Canonization Perspective on Invariant and Equivariant Learning
George Ma*, Yifei Wang*, Derek Lim, Stefanie Jegelka, and Yisen Wang
In NeurIPS, 2024
-
On the Role of Attention Masks and LayerNorm in Transformers
Xinyi Wu, Amir Ajorlou, Yifei Wang, Stefanie Jegelka, and Ali Jadbabaie
In NeurIPS, 2024
-
Dissecting the Failure of Invariant Learning on Graphs
Qixun Wang, Yifei Wang, Yisen Wang, and Xianghua Ying
In NeurIPS, 2024
-
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
-
The Multi-faceted Monosemanticity in Multimodal Representations
Hanqi Yan, Yulan He, and Yifei Wang (Corresponding Author)
In NeurIPS 2024 Workshop on Responsibly Building the Next Generation of Multimodal Foundational Models, 2024
-
Encourage or Inhibit Monosemanticity? Revisit Monosemanticity from a Feature Decorrelation Perspective
Hanqi Yan, Yanzheng Xiang, Guangyi Chen, Yifei Wang, Lin Gui, and Yulan He
In EMNLP, 2024
-
Look Ahead or Look Around? A Theoretical Comparison Between Autoregressive and Masked Pretraining
Qi Zhang, Tianqi Du, Haotian Huang, Yifei Wang, and Yisen Wang
In ICML, 2024
-
OODRobustBench: a benchmark and large-scale analysis of adversarial robustness under distribution shift
Lin Li, Yifei Wang, Chawin Sitawarin, and Michael W. Spratling
In ICML, 2024
-
On the Duality Between Sharpness-Aware Minimization and Adversarial Training
Yihao Zhang, Hangzhou He, Jingyu Zhu, Huanran Chen, Yifei Wang, and Zeming Wei
In ICML, 2024
-
Non-negative Contrastive Learning
Yifei Wang*, Qi Zhang*, Yaoyu Guo, and Yisen Wang
In ICLR, 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 Role of Discrete Tokenization in Visual Representation Learning
Tianqi Du*, Yifei Wang*, and Yisen Wang
In ICLR Spotlight, 2024
2023
-
arXiv Featured by Anthropic Jailbreak and guard aligned language models with only few in-context demonstrations
Zeming Wei, Yifei Wang, and Yisen Wang
arXiv preprint arXiv:2310.06387, 2023
Cited over 250 times. Featured and scaled up in
Anthropic’s blog 📰, where in-context attack successfully jailbroke prominent LLMs including GPT and Claude.
-
Balance, Imbalance, and Rebalance: Understanding Robust Overfitting from a Minimax Game Perspective
Yifei Wang*, Liangchen Li*, Jiansheng Yang, Zhouchen Lin, and Yisen Wang
In NeurIPS, 2023
-
Adversarial Examples Are Not Real Features
Ang Li*, Yifei Wang*, Yiwen Guo, and Yisen Wang
In NeurIPS, 2023
-
Architecture Matters: Uncovering Implicit Mechanisms in Graph Contrastive Learning
Xiaojun Guo*, Yifei Wang*, Zeming Wei, and Yisen Wang
In NeurIPS, 2023
-
Identifiable Contrastive Learning with Automatic Feature Importance Discovery
Qi Zhang*, Yifei Wang*, and Yisen Wang
In NeurIPS, 2023
-
Laplacian Canonization: A Minimalist Approach to Sign and Basis Invariant Spectral Embedding
George Ma*, Yifei Wang*, and Yisen Wang
In NeurIPS, 2023
-
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.
-
Rethinking Weak Supervision in Helping Contrastive Representation Learning
Jingyi Cui*, Weiran Huang*, Yifei Wang*, and Yisen Wang
In ICML, 2023
-
CFA: Class-wise Calibrated Fair Adversarial Training
Zeming Wei, Yifei Wang, Yiwen Guo, and Yisen Wang
In CVPR, 2023
-
Equilibrium Image Denoising with Implicit Differentiation
Qi Chen, Yifei Wang, Zhengyang Geng, Yisen Wang, Jiansheng Yang, and Zhouchen Lin
IEEE Transactions on Image Processing (IEEE TIP), 2023
-
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.
-
ContraNorm: A Contrastive Learning Perspective on Oversmoothing and Beyond
Xiaojun Guo*, Yifei Wang*, Tianqi Du*, and Yisen Wang
In ICLR, 2023
-
Unbiased Stochastic Proximal Solver for Graph Neural Networks with Equilibrium States
Mingjie Li, Yifei Wang, Yisen Wang, and Zhouchen Lin
In ICLR, 2023
-
On the Connection between Invariant Learning and Adversarial Training for Out-of-Distribution Generalization
Shiji Xin, Yifei Wang, Jingtong Su, and Yisen Wang
In AAAI, 2023
2022
-
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.
-
Improving Out-of-distribution Robustness by Adversarial Training with Structured Priors
Qixun Wang*, Yifei Wang*, Hong Zhu, and Yisen Wang
In NeurIPS Spotlight (Top 5%), 2022
-
When Adversarial Training Meets Vision Transformers: Recipes from Training to Architecture
Yichuan Mo, Dongxian Wu, Yifei Wang, Yiwen Guo, and Yisen Wang
In NeurIPS Spotlight (Top 5%), 2022
-
AggNCE: Asymptotically Identifiable Contrastive Learning
Jingyi Cui*, Weiran Huang*, Yifei Wang, and Yisen Wang
In NeurIPS 2022 Self-supervised Learning Workshop (Oral Representation), 2022
-
Optimization-Induced Graph Implicit Nonlinear Diffusion
Qi Chen, Yifei Wang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin
In ICML, 2022
-
G2CN: Graph Gaussian Convolution Networks with Concentrated Graph Filters
Mingjie Li, Xiaojun Guo, Yifei Wang, Yisen Wang, and Zhouchen Lin
In ICML, 2022
-
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
Cited over 130 times. We derived tight generalization bounds for contrastive learning with a new realistic theoretical framework. It derived unsupervised evaluation metrics with 97% correlation to downstream performance.
-
ICLR Silver Best Paper
at ICML-W’21 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.
2021
-
Residual Relaxation for Multi-view Representation Learning
Yifei Wang, Zhengyang Geng, Feng Jiang, Chuming Li, Yisen Wang, Jiansheng Yang, and Zhouchen Lin
In NeurIPS, 2021
-
Dissecting the Diffusion Process in Linear Graph Convolutional Networks
Yifei Wang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin
In NeurIPS, 2021
-
ECML-PKDD Best ML Paper Award 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.