research
My research spans the theory and methodology of the following core areas of modern machine learning (click to see):
- I. Contextual Understanding and Reasoning with LLMs
- II. Unsupervised Representation Learning
- III. Robust Representation Learning
- IV. Structural Representation Learning
Research Methodology. Although people often find that deep learning relies on tons of empirical intuitions, I’d rather believe that “nothing is more intuitive than a good theory.” The process of formalizing a theory of your intuition helps question, sharpen, validate your understanding. I follow this principle in my work and hope you find it helpful as well.
I. Contextual Understanding and Reasoning with LLMs
How LLMs understand, adapt to, and reason with contexts.
I.1 Contextual Understanding: In-context Learning, Long-context Modeling, Length Generalization
- What is Wrong with Perplexity for Long-context Language Modeling? (ICLR 2025)
- Rethinking Invariance in In-context Learning (ICLR 2025)
- Can In-context Learning Really Generalize to OOD Tasks? (ICLR 2025)
- Output Alignment: A Fresh Perspective on Length Generalization in LLMs (ICML 2025)
I.2 Chain-of-thought and Reasoning: Self-correction, Optimal Chain-of-thought Length, Reasoning on Graph
- A Theoretical Understanding of Self-Correction through In-context Alignment (NeurIPS 2024, Best Paper at ICML 2024 Workshop)
- When More is Less: Understanding Chain-of-Thought Length in LLMs (Best Paper Runner-up at ICLR 2025 Workshop)
- G1: Teaching LLMs to Reason on Graphs with Reinforcement Learning (arxiv 2025)
I.3 Transformers: Position Bias, Multi-layer Attention, Dimensional Collapse
- On the Emergence of Position Bias in Transformers (ICML 2025)
- On the Role of Attention Masks and LayerNorm in Transformers (NeurIPS 2024)
- ContraNorm: A Contrastive Learning Perspective on Oversmoothing and Beyond (ICLR 2023)
II. Unsupervised Representation Learning
How to pretrain powerful foundation models from massive unlabeled data.
II.1 Self-supervised Learning (SSL): Contrastive, Non-contrastive, Equivariant, Contextual
- Chaos is a Ladder: A New Theoretical Understanding of Contrastive Learning via Augmentation Overlap (ICLR 2022)
- Towards a Unified Theoretical Understanding of Non-contrastive Learning via Rank Differential Mechanism (ICLR 2023)
- Understanding the Role of Equivariance in Self-supervised Learning (NeurIPS 2024)
- In-Context Symmetries: Self-Supervised Learning through Contextual World Models (NeurIPS 2024)
II.2 Generative Models: Masked Autoencoders, Autoregressive, Energy-based Models, Generative Adversarial Networks
- How Mask Matters: Towards Theoretical Understandings of Masked Autoencoders (NeurIPS 2022)
- Look Ahead or Look Around? A Theoretical Comparison Between Autoregressive and Masked Pretraining (ICML 2024)
- A Unified Contrastive Energy-based Model for Understanding the Generative Ability of Adversarial Training (ICLR 2022)
- Reparameterized Sampling for Generative Adversarial Networks (ECML-PKDD 2021 Best ML Paper)
II.3 Key SSL Components: Predictor, Discrete Tokenization, Projector, Generated Data
- Towards a Unified Theoretical Understanding of Non-contrastive Learning via Rank Differential Mechanism (ICLR 2023)
- Do Generated Data Always Help Contrastive Learning? (ICLR 2024)
- On the Role of Discrete Tokenization in Visual Representation Learning (ICLR 2024)
- Projection Head is Secretly an Information Bottleneck (ICLR 2025)
II.4 Feature Sparsity, Identifiability, Interpretability: Non-negative CL, triCL, Robustness Gains, CSR Embedding
- Non-negative Contrastive Learning (ICLR 2024)
- Identifiable Contrastive Learning with Automatic Feature Importance Discovery (NeurIPS 2023)
- Beyond Interpretability: The Gains of Feature Monosemanticity on Model Robustness (ICLR 2025)
- Beyond Matryoshka: Revisiting Sparse Coding for Adaptive Representation (ICML 2025)
- The Multi-faceted Monosemanticity in Multimodal Representations (NeurIPS Workshop 2024)
III. Robust Representation Learning
How to build models robust to adversarial attacks and reliable across distribution shifts.
III.1 Adversarial Attack and Defense: Adversarial Examples, Unsupervised Adversarial Training, Robust Overfitting
- Adversarial Examples Are Not Real Features (NeurIPS 2023)
- Rethinking the Effect of Data Augmentation in Adversarial Contrastive Learning (ICLR 2023)
- Balance, Imbalance, and Rebalance: Understanding Robust Overfitting from a Minimax Game Perspective (NeurIPS 2023)
III.2 LLM Jailbreak: In-context Attack and Defense, Safety of Chain-of-thought Reasoning
- Jailbreak and guard aligned language models with only few in-context demonstrations (arxiv 2023, 250+ citations)
- Are Smarter LLMs Safer? Exploring Safety-Reasoning Trade-offs in Prompting and Fine-Tuning (arxiv 2025)
III.3 Out-of-distribution (OOD) Generalization: OODRobustBench, Adversarial Training for OOD
- OODRobustBench: Benchmarking Adversarial Robustness under Distribution Shift (ICML 2024)
- Improving Out-of-distribution Robustness by Adversarial Training with Structured Priors (NeurIPS 2022)
- On the Connection between Invariant Learning and Adversarial Training for Out-of-Distribution Generalization (AAAI 2023 Oral)
IV. Structural Representation Learning
How to learn structured data (e.g., graphs) efficiently with structured models, such as, Graph Neural Networks (GNNs).
IV.1 Feature Dynamics of GNNs: oversmoothing, graph equilibrium, unbiased graph sampling, GraphSSL
- Dissecting the Diffusion Process in Linear Graph Convolutional Networks (NeurIPS 2021)
- Optimization-Induced Graph Implicit Nonlinear Diffusion (ICML 2022)
- Unbiased Stochastic Proximal Solver for Graph Neural Networks with Equilibrium States (ICLR 2023)
- Architecture Matters: Uncovering Implicit Mechanisms in Graph Contrastive Learning (NeurIPS 2023)
IV.2 Learning with Symmetry: Laplacian canonicalization, theory of canonicalization