Discovering the Gems in Early Layers: Accelerating Long-Context LLMs with 1000x Input Token Reduction
ACL 2026 Findings
About
I am a Member of Technical Staff at SpaceXAI on the coding specialist team — the core post-training coding team for Grok 4.5. With five teammates, our scope covers webdev, terminal use, kernel, infra, CUA, MLE specialist training, and full-stack long-horizon agentic training. I previously worked on pretraining and agentic search, and have contributed to Grok 4.1, Grok 4.2, Grok 4.3, and Grok 4.5, as well as Grok Build 0.1.
Previously, I was a Senior Research Scientist at MongoDB + Voyage AI, working with Tengyu Ma. I received my Ph.D. in Computer Sciences from the University of Wisconsin–Madison in 2024, advised by Yingyu Liang, and my B.S. in Computer Science and Pure Mathematics Advanced from HKUST in 2019.
Full publication list, awards, and earlier roles are in my CV and PhD thesis.
Publications Click to expand · 29 papers Click to collapse
Kernel Regression in Structured Non-IID Settings: Theory and Implications for Denoising Score Learning
NeurIPS 2025
Unraveling the Smoothness Properties of Diffusion Models: A Gaussian Mixture Perspective
ICCV 2025
Dissecting Submission Limit in Desk-Rejections: A Mathematical Analysis of Fairness in AI Conference Policies
ICML 2025
Fundamental Limits of Visual Autoregressive Transformers: Universal Approximation Abilities
ICML 2025
Beyond Linear Approximations: A Novel Pruning Approach for Attention Matrix
ICLR 2025
When Can We Solve the Weighted Low Rank Approximation Problem in Truly Subquadratic Time?
AIStats 2025
Fourier Circuits in Neural Networks and Transformers: A Case Study of Modular Arithmetic with Multiple Inputs
AIStats 2025
Looped ReLU MLPs May Be All You Need as Practical Programmable Computers
AIStats 2025
Bypassing the Exponential Dependency: Looped Transformers Efficiently Learn In-context by Multi-step Gradient Descent
AIStats 2025
The Computational Limits of State-Space Models and Mamba via the Lens of Circuit Complexity
CPAL 2025 Oral
Fast John Ellipsoid Computation with Differential Privacy Optimization
CPAL 2025 Oral
Curse of Attention: A Kernel-Based Perspective for Why Transformers Fail to Generalize on Time Series Forecasting and Beyond
CPAL 2025
Is A Picture Worth A Thousand Words? Delving Into Spatial Reasoning for Vision Language Models
NeurIPS 2024
Why Larger Language Models Do In-context Learning Differently?
ICML 2024
Towards Few-Shot Adaptation of Foundation Models via Multitask Finetuning
ICLR 2024
Domain Generalization via Nuclear Norm Regularization
CPAL 2024 Oral
Provable Guarantees for Neural Networks via Gradient Feature Learning
NeurIPS 2023
A Graph-Theoretic Framework for Understanding Open-World Semi-Supervised Learning
NeurIPS 2023 Spotlight
When and How Does Known Class Help Discover Unknown Ones? Provable Understandings Through Spectral Analysis
ICML 2023
The Trade-off between Universality and Label Efficiency of Representations from Contrastive Learning
ICLR 2023 Spotlight (Accept Rate: 7.95%)
Experience
Top roles · more below
Member of Technical Staff
Coding specialist team for Grok 4.5 — webdev, terminal use, kernel, infra, CUA, MLE specialist training, and full-stack long-horizon agentic training; previously pretraining and agentic search. Contributed to Grok 4.1, Grok 4.2, Grok 4.3, and Grok 4.5, as well as Grok Build 0.1.
Senior Research Scientist
Embedding & reranking; led rerank-2.5 end-to-end training.
Show more experience Show less
AI Research Scientist Intern
Long-context LLMs, in-context learning, and inference acceleration.
Academic Service
Conference reviewer: ICLR 2022–2026, NeurIPS 2022–2025, ICML 2022 & 2024–2026, COLM 2025, AISTATS 2025–2026, AAAI 2026, IJCAI 2025, CVPR 2021–2022 & 2025–2026, ICCV 2021–2025, ECCV 2020–2022, WACV 2022 & 2025–2026.
Journal reviewer: JVCI, IEEE Transactions on Information Theory.
Recognition: ICML 2026 Gold Reviewer; ICLR 2025 Notable Reviewer.