Yifei Shen

About

I received the B.S. degree in computer science from ShanghaiTech University and Ph.D degree in the electronic and computer engineering from the Hong Kong University of Science and Technology. My research focuses on interpreting the mechanisms of foundation models, analyzing them mathematically, and applying these mechanisms to push the limits of foundation models. Our work consists of three main components:

Controlled Experiments: We design rigorous controlled experiments to train smaller models from scratch, aiming to discover universal laws that extend beyond the current foundation models.

• NeurIPS’24 | ALPINE: Unveiling The Planning Capability of Autoregressive Learning in Language Models (opens in new tab)

Mechanistic Interpretations: We perform mechanistic interpretations of pretrained large models, transforming them into “grey boxes” to enhance our understanding of their inner workings.

• [2411.14982] Large Multi-modal Models Can Interpret Features in Large Multi-modal Models (opens in new tab)
• ICLR’25 | How Do Large Language Models Understand Graph Patterns? A Benchmark for Graph Pattern Comprehension (opens in new tab)
• NeurIPS’24 | Understanding and Improving Training-free Loss-based Diffusion Guidance (opens in new tab)
• ICLR’23 | SIMPLE: Specialized Model-Sample Matching for Domain Generalization (opens in new tab)
• ICLR’23 (Oral) | Sparse Mixture-of-Experts are Domain Generalizable Learners (opens in new tab)

Applications: We use the discovered laws and interpretations to develop novel training paradigms for foundation models. Currently, our applications include interdisciplinary collaborative projects in

• LLM reasoning:
  • NeurIPS’24 | Can Graph Learning Improve Planning in LLM-based Agents? (opens in new tab)
• Embodied AI:
  • [2502.06401] Habitizing Diffusion Planning for Efficient and Effective Decision Making (opens in new tab)
  • ICLR’25 (Spotlight) | What Makes a Good Diffusion Planner for Decision Making? (opens in new tab)
  • [2312.05822] Toward Open-ended Embodied Tasks Solving (opens in new tab)
• AI for Science:
  • [2502.03499] Omni-DNA: A Unified Genomic Foundation Model for Cross-Modal and Multi-Task Learning (opens in new tab)
  • ICLR’24 | Training-free Multi-objective Diffusion Model for 3D Molecule Generation | OpenReview (opens in new tab)
• AI for Communication|Networking|Sensing:
  • IEEE COMMAG’24 | Large Language Models Empowered Autonomous Edge AI for Connected Intelligence (opens in new tab)
  • IEEE JSTSP’23 | An Adaptive and Robust Deep Learning Framework for THz Ultra-Massive MIMO Channel Estimation (opens in new tab)
  • IEEE JSAC’20 | Graph Neural Networks for Scalable Radio Resource Management: Architecture Design and Theoretical Analysis (opens in new tab)

In addition to my primary research, I have a strong interest in system-level work that enables LLM training on less powerful GPUs:

• Full parameter finetuning 8B models on RTX3090 and 70B models on 4 A100s (opens in new tab)

This complementary work aims to make LLMs accessible to more researchers and practitioners.