publication

* denotes shared first authorship

See full list on Google Scholar

2025

  1. arXiv
    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
    PDF Code
  2. ICML Oral
    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
    PDF Code
  3. ICML
    On the Emergence of Position Bias in Transformers
    Xinyi Wu,  Yifei Wang, Stefanie Jegelka, and Ali Jadbabaie
    ICML, 2025
    Featured by MIT News 📰.
    PDF Video
  4. ICML
    Long-Short Alignment for Effective Long-Context Modeling in LLMs
    Tianqi Du*, Haotian Huang*,  Yifei Wang, and Yisen Wang
    ICML, 2025
    PDF Code
  5. 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
    PDF
  6. 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
    PDF Video Code
  7. ICLR LLM training
    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.
    PDF
  8. ICLR
    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.
    PDF
  9. ICLR
    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.
    PDF
  10. ICLR
    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).
    PDF

2024

  1. 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.
    PDF Code
  2. NeurIPS
    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.
    PDF Code
  3. 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 CSAIL News 📰
    We introduced unsupervised test-time adaptation ability to self-supervised learning through a contextual world model designed for joint embedding (JEPA) models.
    PDF Blog Code
  4. NeurIPS
    A Canonization Perspective on Invariant and Equivariant Learning
    George Ma*,  Yifei Wang*, Derek Lim, Stefanie Jegelka, and Yisen Wang
    In NeurIPS, 2024
    PDF Code
  5. NeurIPS
    On the Role of Attention Masks and LayerNorm in Transformers
    Xinyi Wu, Amir Ajorlou,  Yifei Wang, Stefanie Jegelka, and Ali Jadbabaie
    In NeurIPS, 2024
    PDF
  6. NeurIPS
    Dissecting the Failure of Invariant Learning on Graphs
    Qixun Wang,  Yifei Wang, Yisen Wang, and Xianghua Ying
    In NeurIPS, 2024
    PDF Code
  7. NeurIPS Workshop
    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
    PDF
  8. NeurIPS Workshop
    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
    PDF
  9. EMNLP
    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
    PDF Code
  10. ICML
    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
    PDF Code
  11. ICML
    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
    PDF Code Website
  12. ICML
    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
    PDF Code
  13. ICLR
    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.
    PDF Code Slides
  14. ICLR
    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.
    PDF Code
  15. ICLR Spotlight
    On the Role of Discrete Tokenization in Visual Representation Learning
    Tianqi Du*,  Yifei Wang*, and Yisen Wang
    In ICLR Spotlight, 2024
    PDF Code

2023

  1. 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.
    PDF
  2. NeurIPS
    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
    PDF Code
  3. NeurIPS
    Adversarial Examples Are Not Real Features
    Ang Li*,  Yifei Wang*, Yiwen Guo, and Yisen Wang
    In NeurIPS, 2023
    PDF Code
  4. NeurIPS
    Architecture Matters: Uncovering Implicit Mechanisms in Graph Contrastive Learning
    Xiaojun Guo*,  Yifei Wang*, Zeming Wei, and Yisen Wang
    In NeurIPS, 2023
    PDF Code
  5. NeurIPS
    Identifiable Contrastive Learning with Automatic Feature Importance Discovery
    Qi Zhang*,  Yifei Wang*, and Yisen Wang
    In NeurIPS, 2023
    PDF Code
  6. NeurIPS
    Laplacian Canonization: A Minimalist Approach to Sign and Basis Invariant Spectral Embedding
    George Ma*,  Yifei Wang*, and Yisen Wang
    In NeurIPS, 2023
    PDF Code
  7. ICML
    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.
    PDF Code
  8. ICML
    Rethinking Weak Supervision in Helping Contrastive Representation Learning
    Jingyi Cui*, Weiran Huang*,  Yifei Wang*, and Yisen Wang
    In ICML, 2023
    PDF
  9. CVPR
    CFA: Class-wise Calibrated Fair Adversarial Training
    Zeming Wei,  Yifei Wang, Yiwen Guo, and Yisen Wang
    In CVPR, 2023
    PDF Code
  10. TIP
    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
    PDF
  11. ICLR
    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.
    PDF Blog Code Slides
  12. ICLR
    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.
    PDF Code
  13. ICLR
    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.
    PDF Code
  14. ICLR
    ContraNorm: A Contrastive Learning Perspective on Oversmoothing and Beyond
    Xiaojun Guo*,  Yifei Wang*, Tianqi Du*, and Yisen Wang
    In ICLR, 2023
    PDF Code
  15. ICLR
    Unbiased Stochastic Proximal Solver for Graph Neural Networks with Equilibrium States
    Mingjie Li,  Yifei Wang, Yisen Wang, and Zhouchen Lin
    In ICLR, 2023
    PDF
  16. AAAI Oral
    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
    PDF

2022

  1. NeurIPS Spotlight
    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.
    PDF Code Slides
  2. NeurIPS Spotlight
    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
    PDF Code Slides
  3. NeurIPS Spotlight
    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
    PDF Code
  4. NeurIPS Workshop Oral
    AggNCE: Asymptotically Identifiable Contrastive Learning
    Jingyi Cui*, Weiran Huang*,  Yifei Wang, and Yisen Wang
    In NeurIPS 2022 Self-supervised Learning Workshop (Oral Representation), 2022
    PDF
  5. ICML
    Optimization-Induced Graph Implicit Nonlinear Diffusion
    Qi Chen,  Yifei Wang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin
    In ICML, 2022
    PDF Code
  6. ICML
    G2CN: Graph Gaussian Convolution Networks with Concentrated Graph Filters
    Mingjie Li, Xiaojun Guo,  Yifei Wang, Yisen Wang, and Zhouchen Lin
    In ICML, 2022
    PDF
  7. ICLR
    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.
    PDF Code Slides
  8. 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.
    PDF Slides

2021

  1. NeurIPS
    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
    PDF Blog Slides
  2. NeurIPS
    Dissecting the Diffusion Process in Linear Graph Convolutional Networks
    Yifei Wang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin
    In NeurIPS, 2021
    PDF Blog Code Slides
  3. 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.
    PDF Code Slides