Hello! I am a postdoctoral researcher at MIT CSAIL working with Stefanie Jegelka.
I obtained my PhD in Applied Mathematics from Peking University (PKU) in 2023, advised by Yisen Wang, Jiansheng Yang, and Zhouchen Lin. Prior to that, I got my bachelor's degrees from PKU math and phil.
My first-author works received Best ML Paper Award of ECML-PKDD 2021, Silver Best Paper Award of ICML 2021 AdvML workshop, and Spotlight Award at ICML 2024 ICL workshop.
I am generally interested in machine learning and representation learning, with a focus on the following topics:
- Self-supervised Learning (SSL). SSL is the driving engine for pretraining foundation models. I am mostly interested in theoretical understanding and principled design of SSL (including generative and non-generative models), aimed at uncovering their inherent mechanisms and enhancing model interpretability.
- Safe Alignment. Powerful LLMs need to be aligned to human purposes with guardrails to avoid being abused. I explore how to systematically develop robust foundation models against real-world distribution shifts.
- Neural Representations. I am interested in understanding the inherent representations and feature dynamics of general-purpose neural network backbones, in particular, Transformers and Graph Neural Networks.
Contact: yifei_w at mit.edu / Google Scholar / Github / X (Twitter)
News
- 2024.09. π 6 papers were accepted to NeurIPS 2024, covering LLM self-correction, In-context SSL, Equivariant SSL, Canonization-based Equivariance, Transformer Dynamics, and Invariant Learning on Graph.
- 2024.09. 1 paper was accepted to EMNLP 2024.
- 2024.09. I gave a talk at NYU Tandon on Building Safe Foundation Models from Principled Understanding.
- 2024.08. I will be co-organizing ML Tea seminar at MIT CSAIL this fall. Join us on the Mondays at 32-G882!
- 2024.08. I gave a talk at Princeton University on Reimagining Self-Supervised Learning with Context.
- 2024.08. I wll be serving as an Area Chair for ICLR 2025.
- 2024.07. Our paper on theoretical understanding of LLM self-correction received the Spotlight Award (awarded to top 3 papers) at ICML 2024 1st ICL Workshop.
- 2024.07. I will be co-organizing the NeurIPS 2024 workshop on Red Teaming GenAI.
- 2024.07. 3 papers were accepted by ICML 2024.
- 2024.06. I gave a talk on Non-negative Contrastive Learning at Cohere AI.
- 2024.05. I visited TU Munich and gave a talk on Self-supervised Learning of Identifiable Features.
- 2024.05. I served as a Session Chair at ICLR 2024.
- 2024.04. I gave a talk at MIT LIDS Tea on Non-negative Contrastive Learning (slides).
- 2024.02. I was honored to receive Wenjun Wu Outstanding Ph.D. Dissertation Runner-Up Award (top 14 nation-wide) from CAAI. Wenjun Wu invented Wu's method for automatic theorem proving and pioneered AI research in China. Thanks everyone!
- 2024.01. 3 papers were accepted to ICLR 2024.
- 2023.12. I have joined MIT CSAIL as a postdoc.
- 2023.09. I will be serving as an Area Chair for ICLR 2024.
Papers (* marks equal contribution)
-
A Theoretical Understanding of Self-Correction through In-context Alignment
NeurIPS 2024
ICML 2024 Workshop on In-Context Learning (ICL) Spotlight Award (awarded to top 3 papers) 2024 PDF - Understanding the Role of Equivariance in Self-supervised Learning NeurIPS 2024 2024 PDF
- In-Context Symmetries: Self-Supervised Learning through Contextual World Models NeurIPS 2024 2024 PDF
- A Canonization Perspective on Invariant and Equivariant Learning NeurIPS 2024 2024 PDF
- On the Role of Attention Masks and LayerNorm in Transformers NeurIPS 2024 2024 PDF
- Dissecting the Failure of Invariant Learning on Graphs NeurIPS 2024 2024 PDF
- Encourage or Inhibit Monosemanticity? Revisit Monosemanticity from a Feature Decorrelation Perspective EMNLP 2024 2024 PDF
-
A Theoretical Understanding of Self-Correction through In-context Alignment
NeurIPS 2024
ICML 2024 Workshop on In-Context Learning (ICL) Spotlight Award (awarded to top 3 papers) 2024 PDF - Understanding the Role of Equivariance in Self-supervised Learning NeurIPS 2024 2024 PDF
- In-Context Symmetries: Self-Supervised Learning through Contextual World Models NeurIPS 2024 2024 PDF
- A Canonization Perspective on Invariant and Equivariant Learning NeurIPS 2024 2024 PDF
- On the Role of Attention Masks and LayerNorm in Transformers NeurIPS 2024 2024 PDF
- Dissecting the Failure of Invariant Learning on Graphs NeurIPS 2024 2024 PDF
- Encourage or Inhibit Monosemanticity? Revisit Monosemanticity from a Feature Decorrelation Perspective EMNLP 2024 2024 PDF
- Look Ahead or Look Around? A Theoretical Comparison Between Autoregressive and Masked Pretraining ICML 2024 2024 PDF | Code
- OODRobustBench: a benchmark and large-scale analysis of adversarial robustness under distribution shift ICML 2024 2024 PDF | Code | Leaderboard
- On the Duality Between Sharpness-Aware Minimization and Adversarial Training ICML 2024 2024 PDF | Code
- Rethinking Invariance in In-context Learning ICML 2024 Workshop on Theoretical Foundations of Foundation Models (TF2M) 2024 PDF
- Non-negative Contrastive Learning ICLR 2024 2024 PDF | Code | Slides
-
Do Generated Data Always Help Contrastive Learning?
ICLR 2024
2024
PDF |
Code |
Featured on Sync (CN)
- On the Role of Discrete Tokenization in Visual Representation Learning ICLR 2024 (Spotlight) 2024 PDF | Code
- How to Craft Backdoors with Unlabeled Data Alone? ICLR 2024 Workshop on Navigating and Addressing Data Problems for Foundation Models (DPFM) 2024 PDF
- Virtual Classifier: A Reversed Approach for Robust Image Evaluation ICLR 2024 Workshop on Navigating and Addressing Data Problems for Foundation Models (DPFM) 2024 PDF
- Balance, Imbalance, and Rebalance: Understanding Robust Overfitting from a Minimax Game Perspective NeurIPS 2023 2023 PDF | Code
- Adversarial Examples Are Not Real Features NeurIPS 2023 2023 PDF | Code
- Architecture Matters: Uncovering Implicit Mechanisms in Graph Contrastive Learning NeurIPS 2023 2023 PDF | Code
- Identifiable Contrastive Learning with Automatic Feature Importance Discovery NeurIPS 2023 2023 PDF | Code
- Laplacian Canonization: A Minimalist Approach to Sign and Basis Invariant Spectral Embedding NeurIPS 2023 2023 PDF | Code
- On the Generalization of Multi-modal Contrastive Learning ICML 2023 2023 PDF | Code
- Rethinking Weak Supervision in Helping Contrastive Representation Learning ICML 2023 2023 PDF
- CFA: Class-wise Calibrated Fair Adversarial Training CVPR 2023 2023 PDF | Code
- Equilibrium Image Denoising with Implicit Differentiation IEEE Transactions on Image Processing (TIP) 2023 PDF
- A Message Passing Perspective on Learning Dynamics of Contrastive Learning ICLR 2023 2023 PDF | Code | Slides | Blog
- Towards a Unified Theoretical Understanding of Non-contrastive Learning via Rank Differential Mechanism ICLR 2023 2023 PDF | Code
- Rethinking the Effect of Data Augmentation in Adversarial Contrastive Learning ICLR 2023 2023 PDF | Code
- ContraNorm: A Contrastive Learning Perspective on Oversmoothing and Beyond ICLR 2023 2023 PDF | Code
- Unbiased Stochastic Proximal Solver for Graph Neural Networks with Equilibrium States ICLR 2023 2023 PDF
- What Contrastive Learning Learns Beyond Class-wise Features? ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models (ME-FoMo) 2023 PDF
- Rethinking the Necessity of Labels in Backdoor Defense ICLR 2023 Workshop on Backdoor Attacks and Defenses in Machine Learning (BANDS) 2023 PDF
- On the Connection between Invariant Learning and Adversarial Training for Out-of-Distribution Generalization AAAI 2023 (Oral) 2023 PDF
- How Mask Matters: Towards Theoretical Understandings of Masked Autoencoders NeurIPS 2022 (Spotlight, Top 5%) 2022 PDF | Code | Slides
- Improving Out-of-distribution Robustness by Adversarial Training with Structured Priors NeurIPS 2022 (Spotlight, Top 5%) 2022 PDF | Code | Slides
- When Adversarial Training Meets Vision Transformers: Recipes from Training to Architecture NeurIPS 2022 (Spotlight, Top 5%) 2022 PDF | Code
- Variational Energy-Based Models: A Probabilistic Framework for Contrastive Self-Supervised Learning NeurIPS 2022 SSL Workshop 2022 PDF
- AggNCE: Asymptotically Identifiable Contrastive Learning NeurIPS 2022 SSL Workshop (Oral) 2022 PDF
- Efficient and Scalable Implicit Graph Neural Networks with Virtual Equilibrium IEEE BigData 2022 (Long Talk) 2022 PDF
- Optimization-Induced Graph Implicit Nonlinear Diffusion ICML 2022 2022 PDF | Code
- G2CN: Graph Gaussian Convolution Networks with Concentrated Graph Filters ICML 2022 2022 PDF
- Chaos is a Ladder: A New Theoretical Understanding of Contrastive Learning via Augmentation Overlap ICLR 2022 2022 PDF | Code | Slides
- A Unified Contrastive Energy-based Model for Understanding the Generative Ability of Adversarial Training ICLR 2022 (π Silver Best Paper Award @ ICML 2021 AdvML workshop) 2022 PDF | Slides | Award
- Residual Relaxation for Multi-view Representation Learning NeurIPS 2021 2021 PDF | Slides | Blog
- Dissecting the Diffusion Process in Linear Graph Convolutional Networks NeurIPS 2021 2021 PDF | Code | Slides | Blog
- Reparameterized Sampling for Generative Adversarial Networks ECML-PKDD 2021 2021 (π Best ML Paper Award (1/685) & Invited to Machine Learning Journal) PDF | Code | Slides | Media | Talk | Award
- Train Once, and Decode as You Like COLING 2020 2020 PDF
-
A Theoretical Understanding of Self-Correction through In-context Alignment
NeurIPS 2024
ICML 2024 Workshop on In-Context Learning (ICL) Spotlight Award (awarded to top 3 papers) 2024 PDF - Understanding the Role of Equivariance in Self-supervised Learning NeurIPS 2024 2024 PDF
- In-Context Symmetries: Self-Supervised Learning through Contextual World Models NeurIPS 2024 2024 PDF
- Encourage or Inhibit Monosemanticity? Revisit Monosemanticity from a Feature Decorrelation Perspective EMNLP 2024 2024 PDF
- Look Ahead or Look Around? A Theoretical Comparison Between Autoregressive and Masked Pretraining ICML 2024 2024 PDF | Code
- Rethinking Invariance in In-context Learning ICML 2024 Workshop on Theoretical Foundations of Foundation Models (TF2M) 2024 PDF
- Non-negative Contrastive Learning ICLR 2024 2024 PDF | Code | Slides
- Do Generated Data Always Help Contrastive Learning? ICLR 2024 2024 PDF | Code | Featured on Sync (CN)
- On the Role of Discrete Tokenization in Visual Representation Learning ICLR 2024 (Spotlight) 2024 PDF | Code
- How to Craft Backdoors with Unlabeled Data Alone? ICLR 2024 Workshop on Navigating and Addressing Data Problems for Foundation Models (DPFM) 2024 PDF
- Virtual Classifier: A Reversed Approach for Robust Image Evaluation ICLR 2024 Workshop on Navigating and Addressing Data Problems for Foundation Models (DPFM) 2024 PDF
- Architecture Matters: Uncovering Implicit Mechanisms in Graph Contrastive Learning NeurIPS 2023 2023 PDF | Code
- Identifiable Contrastive Learning with Automatic Feature Importance Discovery NeurIPS 2023 2023 PDF | Code
- On the Generalization of Multi-modal Contrastive Learning ICML 2023 2023 PDF | Code
- Rethinking Weak Supervision in Helping Contrastive Representation Learning ICML 2023 2023 PDF
- A Message Passing Perspective on Learning Dynamics of Contrastive Learning ICLR 2023 2023 PDF | Code | Slides | Blog
- Towards a Unified Theoretical Understanding of Non-contrastive Learning via Rank Differential Mechanism ICLR 2023 2023 PDF | Code
- Rethinking the Effect of Data Augmentation in Adversarial Contrastive Learning ICLR 2023 2023 PDF | Code
- ContraNorm: A Contrastive Learning Perspective on Oversmoothing and Beyond ICLR 2023 2023 PDF | Code
- What Contrastive Learning Learns Beyond Class-wise Features? ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models (ME-FoMo) 2023 PDF
- Rethinking the Necessity of Labels in Backdoor Defense ICLR 2023 Workshop on Backdoor Attacks and Defenses in Machine Learning (BANDS) 2023 PDF
- How Mask Matters: Towards Theoretical Understandings of Masked Autoencoders NeurIPS 2022 (Spotlight, Top 5%) 2022 PDF | Code | Slides
- Variational Energy-Based Models: A Probabilistic Framework for Contrastive Self-Supervised Learning NeurIPS 2022 SSL Workshop 2022 PDF
- AggNCE: Asymptotically Identifiable Contrastive Learning NeurIPS 2022 SSL Workshop (Oral) 2022 PDF
- Chaos is a Ladder: A New Theoretical Understanding of Contrastive Learning via Augmentation Overlap ICLR 2022 2022 PDF | Code | Slides
- A Unified Contrastive Energy-based Model for Understanding the Generative Ability of Adversarial Training ICLR 2022 (π Silver Best Paper Award @ ICML 2021 AdvML workshop) 2022 PDF | Slides | Award
- Residual Relaxation for Multi-view Representation Learning NeurIPS 2021 2021 PDF | Slides | Blog
- Reparameterized Sampling for Generative Adversarial Networks ECML-PKDD 2021 2021 (π Best ML Paper Award (1/685). Invited to Machine Learning Journal) PDF | Code | Slides | Media | Talk | Award
-
A Theoretical Understanding of Self-Correction through In-context Alignment
NeurIPS 2024
ICML 2024 Workshop on In-Context Learning (ICL) Spotlight Award (top 3 papers) 2024 PDF - OODRobustBench: a benchmark and large-scale analysis of adversarial robustness under distribution shift ICML 2024 2024 PDF | Code | Leaderboard
- On the Duality Between Sharpness-Aware Minimization and Adversarial Training ICML 2024 2024 PDF | Code
- How to Craft Backdoors with Unlabeled Data Alone? ICLR 2024 Workshop on Navigating and Addressing Data Problems for Foundation Models (DPFM) 2024 PDF
- Adversarial Examples Are Not Real Features NeurIPS 2023 2023 PDF | Code
- Balance, Imbalance, and Rebalance: Understanding Robust Overfitting from a Minimax Game Perspective NeurIPS 2023 2023 PDF | Code
- CFA: Class-wise Calibrated Fair Adversarial Training CVPR 2023 2023 PDF | Code
- Rethinking the Effect of Data Augmentation in Adversarial Contrastive Learning ICLR 2023 2023 PDF
- Rethinking the Necessity of Labels in Backdoor Defense ICLR 2023 Workshop on Backdoor Attacks and Defenses in Machine Learning (BANDS) 2023 PDF
- On the Connection between Invariant Learning and Adversarial Training for Out-of-Distribution Generalization AAAI 2023 (Oral) 2023 PDF
- Improving Out-of-distribution Robustness by Adversarial Training with Structured Priors NeurIPS 2022 (Spotlight, Top 5%) 2022 PDF | Code | Slides
- When Adversarial Training Meets Vision Transformers: Recipes from Training to Architecture NeurIPS 2022 (Spotlight, Top 5%) 2022 PDF | Code
- A Unified Contrastive Energy-based Model for Understanding the Generative Ability of Adversarial Training ICLR 2022 (π Silver Best Paper Award @ ICML 2021 AdvML workshop) 2022 PDF | Slides | Award
- A Canonization Perspective on Invariant and Equivariant Learning NeurIPS 2024 2024 PDF
- On the Role of Attention Masks and LayerNorm in Transformers NeurIPS 2024 2024 PDF
- Dissecting the Failure of Invariant Learning on Graphs NeurIPS 2024 2024 PDF
- Laplacian Canonization: A Minimalist Approach to Sign and Basis Invariant Spectral Embedding NeurIPS 2023 2023 PDF | Code
- Architecture Matters: Uncovering Implicit Mechanisms in Graph Contrastive Learning NeurIPS 2023 2023 PDF | Code
- ContraNorm: A Contrastive Learning Perspective on Oversmoothing and Beyond ICLR 2023 2023 PDF
- Unbiased Stochastic Proximal Solver for Graph Neural Networks with Equilibrium States ICLR 2023 2023 PDF
- Efficient and Scalable Implicit Graph Neural Networks with Virtual Equilibrium IEEE BigData 2022 (Long Talk) 2022 PDF
- Optimization-Induced Graph Implicit Nonlinear Diffusion ICML 2022 2022 PDF | Code
- G2CN: Graph Gaussian Convolution Networks with Concentrated Graph Filters ICML 2022 2022 PDF
- Dissecting the Diffusion Process in Linear Graph Convolutional Networks NeurIPS 2021 2021 PDF | Code | Slides | Blog
Selected Awards
Spotlight Award (top 3), ICML 2024 ICL workshop, 2024.
Wenjun Wu Outstanding Ph.D. Dissertation Runner-Up Award (top 14 nation-wide), CAAI, 2024
Excellent Graduate (top 1 per department), Beijing, 2023
Excellent Graduate, Peking University, 2023
Baidu Scholarship Runner-Up (top 20 nation-wide), Baidu Inc, 2022
National Scholarship (top 0.1% nation-wide), China, 2021, 2022
Principal Scholarship (top 1% university-wide), Peking University, 2022
Best ML Paper Award (1/685), ECML-PKDD 2021, 2021
Silver Best Paper Award, ICML 2021 AdvML workshop, 2021
Wenjun Wu Outstanding Ph.D. Dissertation Runner-Up Award (top 14 nation-wide), CAAI, 2024
Excellent Graduate (top 1 per department), Beijing, 2023
Excellent Graduate, Peking University, 2023
Baidu Scholarship Runner-Up (top 20 nation-wide), Baidu Inc, 2022
National Scholarship (top 0.1% nation-wide), China, 2021, 2022
Principal Scholarship (top 1% university-wide), Peking University, 2022
Best ML Paper Award (1/685), ECML-PKDD 2021, 2021
Silver Best Paper Award, ICML 2021 AdvML workshop, 2021
Community Services
Reviewer:
- Conferences: NeurIPS (2022, 2023, 2024), ICML (2022), AISTATS (2024, 2025), LoG (2023, 2024), ECML-PKDD (2022), AAAI (2025), CVPR (2023, 2024), ICCV (2023), ACL (2020, 2021)
- Journal: IEEE TPAMI, TMLR, TKDE
- ICLR (2024, 2025)
- NeurIPS 2024 Workshop Red Teaming GenAI: What Can We Learn from Adversaries
- ML Tea Seminar of MIT CSAIL (2024 Fall)