Yifei Wang

avatar.JPG

I am a Member of Technical Staff at Amazon AGI SF Lab. I help build the lab’s next-generation general-agent training effort from scratch, shaping its agent RL system, harness design, data pipeline, and post-training strategy.

I was a postdoc at MIT CSAIL (2023-2025), advised by Stefanie Jegelka. I received my Ph.D. in Applied Mathematics from Peking University, advised by Yisen Wang, Zhouchen Lin, and Jiansheng Yang. I also completed my B.S. and B.A. at Peking University.

In research, I am generally interested in reasoning, representation learning, and self-supervised learning. My work has received 5 best paper awards and has been featured by MIT News and Anthropic. I serve as an Area Chair for ICLR, ICML and NeurIPS.

news

December, 2025 Our paper G1 has received the Best Paper Award at the NeurIPS 2025 NPGML Workshop.
August, 2025 I gave an invited talk Two New Dimensions of Sparsity for Scaling LLMs at Google DeepMind on sparse long-context training (ICLR 2025) and sparse embedding (ICML 2025 Oral).
June, 2025 Our ICML 2025 paper was featured in an MIT News article, Unpacking the bias of large language models, where we identified and theoretically proved the root causes of position bias in Transformers.
June, 2025 I gave an invited talk at the ASAP Seminar on Your Next-Token Prediction and Transformers Are Biased for Long-Context Modeling—see the recording at YouTube.
April, 2025 Our recent work When More is Less: Understanding Chain-of-Thought Length in LLMs received the Best Paper Runner-Up Award 🏆 at ICLR 2025 Workshop on Reasoning and Planning for LLMs.
April, 2025 I will give a tutorial on the Principles of Self-supervised Learning in the Foundation Model Era at IJCAI 2025 (Aug 16 - Aug 22).

selected papers (see full publication)

  1. ICML
    SSL4RL: Revisiting Self-supervised Learning as Intrinsic Reward for Visual-Language Reasoning
    Xiaojun Guo*, Runyu Zhou*,  Yifei Wang*, Qi Zhang, Chenheng Zhang, Stefanie Jegelka, Xiaohan Wang, Jiajun Chai, Guojun Yin, Wei Lin, and Yisen Wang
    ICML, 2026
    PDF
  2. ICLR 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, 2026
    🏆 Best Paper Runner-up Award at ICLR 2025 Workshop on Reasoning and Planning for LLMs
    PDF
  3. NeurIPS Best Paper
    G1: Teaching LLMs to Reason on Graphs with Reinforcement Learning
    Xiaojun Guo*, Ang Li*,  Yifei Wang*, Stefanie Jegelka, and Yisen Wang
    In NeurIPS, 2025
    🏆 Best Paper Award at NeurIPS 2025 NPGML Workshop
    PDF Code
  4. 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 (Top 1%), 2025
    PDF Code
  5. 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
  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. 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
  8. TPAMI Featured by Anthropic
    Jailbreak and guard aligned language models with only few in-context demonstrations
    Zeming Wei,  Yifei Wang, and Yisen Wang
    IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI), arxiv:2310.06387, 2023
    Featured and scaled up by Anthropic to many-shot jailbreaking. Cited over 400 times
    PDF
  9. ICLR JMLR
    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
    A new augmentation overlap theory for understanding the generalization of contrastive learning. Cited over 150 times. Extended version was accepted at JMLR.
    PDF Code Slides