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

April, 2026 Four papers were accepted to ICML 2026, covering self-supervised rewards for RL (SSL4RL), single-stage sparse multi-vector retrieval (SSR), agentic graph reasoning (GRASP), and data scale-difficulty interplay. One latent visual reasoning work, Monet, was accepted to CVPR 2026.
January, 2026 Six papers were accepted to ICLR 2026, explaining scaling laws of CoT length, showing AR models rival diffusion models for any-order generation, and exploring real-world benefits of sparsity via ultra-sparse embeddings, sparse feature attention, and predicting LLM transferability.
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.

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
  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
  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
  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
  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
  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
  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.
  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
  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.