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 AdvML workshop, and the Best Paper Award at ICML 2024 ICL workshop. My thesis won the CAAI Outstanding Ph.D. Dissertation Runner-Up Award. I have published 33 papers at NeurIPS, ICLR, and ICML, and I am a (co-)first author on 22 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

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.
November, 2024 I gave a guest lecture on Towards Test-time Self-supervised Learning (slides) at Boston College.
October, 2024 3 new papers are on arxiv, exploring 1) how existing long-context training of LLMs is problematic and how to address it (paper), 2) how sparse autoencoders can significantly improve robustness at noisy and few-shot scenarios (paper), and 3) whether ICL can truly extrapolate to OOD scenarios (paper).
October, 2024 6 papers were accepted to NeurIPS 2024. We inverstigated how LLMs perform self-correction at test time (paper), how to build dynamic world models through joint embedding methods (paper), how Transformers avoid feature collapse with LayerNorm and attention masks (paper), and why equivariant prediction of data corruptions helps learn good representations (paper).
September, 2024 I gave a talk at NYU Tandon on Building Safe Foundation Models from Principled Understanding.
August, 2024 I gave a talk at Princeton University on Reimagining Self-Supervised Learning with Context.

selected publications

  1. arXiv LLM training & 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
    arXiv preprint arXiv:2410.23771, 2024
    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.