Yifei Wang (on the job market!)

Postdoctoral researcher at MIT CSAIL, advised by Stefanie Jegelka.

avatar.JPG

My goal is to develop models that learn from massive data with minimal human efforts, which drives my persistent interests in self-supervised foundation models. My research has contributed to unveiling the key principles underlying these foundation models and designing efficient algorithms to improve their capabilities and safety:

  • Mathematical Principles of Foundation Models. We established theoretical foundations for a broad spectrum of Self-Supervised Learning (SSL) methods that are at the heart of foundation models, from contrastive [1, 2], non-contrastive [3], autoregressive [4], reconstructive [5], to predictive [6] approaches. Our recent work further pioneered the first rigorous theory [7] for the test-time self-correction ability of LLMs, a key mechanism for scaling reasoning during inference.
  • Improving Model Capabilities. We leveraged these principles to “debug” and “boost” foundation models. We generalized self-supervised learning to be able to self-adapt to new tasks without retraining [8] (featured by MIT), proposed adaptive training with AI data to circumvent data shortage [9], and significantly enhanced LLMs’ long-context understanding through self-identifying key tokens [10].
  • Safe and Trustworthy AI. We developed principled understandings and algorithms for adversarial robustness [11, 12, 13, 14], interpretability [15, 16], and domain generalization [17, 18, 19]. In DynACL [20], we built the first self-supervised model that is as robust as the supervised one. We firstly showed that LLMs’ core emergent abilities, in-context learning [21] and self-correction [7], can play important roles in safety tasks like jailbreaking, which was featured and scaled up by Anthropic.

My first-author papers received the Best ML Paper Award (1/685) at ECML-PKDD 2021, the Silver Best Paper Award at ICML 2021 workshop, and the Best Paper Award at ICML 2024 workshop. My thesis won CAAI Outstanding Ph.D. Dissertation Runner-Up Award. I have published 44 peer-reviewed papers (39 in NeurIPS, ICLR, and ICML), and I am a (co-)first author on 28 of them.

I served as an organizer for NeurIPS 2024 Workshop on Red Teaming GenAI and the ML Tea Seminar at MIT. I served as an Area Chair for ICLR 2024 and 2025, and as a reviewer for main AI conferences (NeurIPS, ICML, ECML, AISTATS, LoG, CVPR, ACL).

I obtained my PhD in Applied Mathematics from Peking University in 2023, advised by Yisen Wang, Zhouchen Lin, Jiansheng Yang. Prior to that, I did my undergraduate at School of Mathematical Sciences, Peking University.

I am on the job market 2024-2025 and actively looking for jobs! Links: CV | Research Statement

news

January, 2025 6 papers were accepted at ICLR 2025 (3 as a co-first author)! We proposed long-context perplexity, invariant in-context learning, constrained tool decoding for better training and usage of LLMs. We also looked into some fundamental questions, such as OOD generalization of in-context learning, interplay between monosemanticity and robustness, and the nature of projection heads.
January, 2025 I will give a talk at the CILVR seminar at NYU CDS on Feb 5.
January, 2025 I will give a talk at Boston University on Jan 29.
December, 2024 Our NeurIPS’24 work ContextSSL was featured by MIT 📰: Machines that Self-adapt to New Tasks without Re-training. It was also selected as an oral presentation (top 4) at NeurIPS’24 SSL workshop.
December, 2024 I gave a talk on Principles of Foundations Models at Johns Hopkins University.

selected publications

  1. 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
    We proposed a long-context perplexity (LongPPL) measure that emphasizes long-context relevant tokens at training and evaluation, improving benchmark scores on LongBench, LongEval, and RULER by up to 22%.
  2. 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.
  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 📰
    We introduced unsupervised test-time adaptation ability to self-supervised learning through a contextual world model designed for joint embedding (JEPA) models.
  4. 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 160 times. Featured and scaled up in Anthropic’s blog 📰, where in-context attack successfully jailbroke prominent LLMs including GPT and Claude.
  5. 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.
  6. 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.
  7. 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.
  8. 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.
  9. 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.
  10. 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.