Xuchen Gong

Hi! I am a second-year Ph.D. student in the Computer Science Department at the University of Chicago advised by Prof. Tian Li. I study how learning systems make decisions when the feedback driving them is weak, delayed, non-verifiable, or non-differentiable. In LLM optimization settings, I worked on gradient estimation under restricted feedback. In agentic and interactive systems, the decisions are which tool to call, which subproblem to pursue, and which agent to collaborate with. I am interested in credit assignment over long trajectories, reward design and data curation for open-ended tasks, and joint learning of policies, rewards, and environments.

Previously, I studied tradeoffs between the primary objective and other properties such as differential privacy and robustness during my M.S. in Machine Learning at Carnegie Mellon University. Before that, I received my B.S. in Data Science from Duke Kunshan University and Duke University.

I am always excited to discuss research perspectives and connect with people. If you are interested, please feel free to send me an email!

Selected Work

* indicates equal contribution

Learning Scientific Exploration from Human Research Decision Trajectories

Xuchen Gong, Shane Gu, Haokun Liu, Dixi Yao, Chenhao Tan, and Tian Li

arXiv 2026

Abstract Paper Blog Code

A key challenge in building AI systems for scientific research is enabling \textitscientific exploration: the systematic process of investigating unknown phenomena or ideas to gain new knowledge through sequences of research decisions and actions. Yet this process is largely missing from existing scientific corpora; for example, research papers primarily record final outcomes rather than the trajectories that produced them. In this work, we introduce ResearchTrails, a dataset of human research trajectories constructed from Git repositories, where commit histories serve as proxies for research exploration. We develop an automated and scalable pipeline that extracts structured research trajectories from repository commits, capturing successive changes to methods, experiments, and ablations. We characterize the resulting dataset and show that these trajectories contain meaningful signals about intermediate research decisions beyond what final papers reveal. We further demonstrate utilities of ResearchTrails in multiple use cases, including retrieving human research experience as external skills at test time and training models on research trajectories to improve generalization to new research decisions. Our results suggest a path toward AI systems that learn not only from the products of science, but from the evolving process of discovery itself.
Principled Top-k Selection for Language Models with Hybrid Gradients

Xuchen Gong, Junfei Sun, and Tian Li

arXiv 2026

Abstract arXiv

Selecting the best k items out of m candidates is a critical component of modern large language model systems, such as document selection in Retrieval-Augmented Generation (RAG) and expert routing in Mixture-of-Experts (MoEs). However, training these selection modules remains challenging due to weak gradient signals and suboptimal exploration-exploitation tradeoffs. Furthermore, prior works often rely on heuristics, lacking principled objectives and approaches that explicitly model and solve the top-k selection problem. In this work, we propose a principled objective for training selection modules, whose gradient naturally provides richer training signals in a hybrid form—containing a supervised-gradient component and a policy-gradient component. We show that the selection problem becomes harder as m increases, and our algorithm converges at rate O(1/\sqrtT), with the optimal upper bound achieved by balancing between bias and variance. Practically, we apply our method to a set of tasks involving top-k selection, including synthetic regression problems, RAG, and MoE systems, showing that our method outperforms the baselines in next-token prediction perplexity and QA accuracies.
Asynchronous Heavy-Tailed Optimization

Junfei Sun, Dixi Yao, Xuchen Gong, Tahseen Rabbani, Manzil Zaheer, and Tian Li

arXiv 2026

Abstract arXiv

Heavy-tailed stochastic gradient noise, commonly observed in transformer models, can destabilize the optimization process. Recent works mainly focus on developing and understanding approaches to address heavy-tailed noise in the centralized or distributed, synchronous setting, leaving the interactions between such noise and asynchronous optimization underexplored. In this work, we investigate two communication schemes that handle stragglers with asynchronous updates in the presence of heavy-tailed gradient noise. We propose and theoretically analyze algorithmic modifications based on delay-aware learning rate scheduling and delay compensation to enhance the performance of asynchronous algorithms. Our convergence guarantees under heavy-tailed noise match the rate of the synchronous counterparts and improve delay tolerance compared with existing asynchronous approaches. Empirically, our approaches outperform prior synchronous and asynchronous methods in terms of accuracy/runtime trade-offs and are more robust to hyperparameters in both image and language tasks.
Zeroth-Order Sharpness-Aware Learning with Exponential Tilting

Xuchen Gong, and Tian Li

NeurIPS 2026

Abstract arXiv Code

Classic zeroth-order optimization approaches typically optimize for a smoothed version of the original function, i.e., the expected objective under randomly perturbed model parameters. This can be interpreted as encouraging the loss values in the perturbation set to be small on average. Popular sharpness-aware minimization (SAM) objectives, however, typically focus on the largest loss within the neighborhood to arrive at flat minima more effectively. In this work, we connect zeroth-order optimization (and its corresponding objectives) with SAM approaches explicitly, through an exponential tilting objective that provides a smooth transition between the average- and the max-loss formulations. We explore new zeroth-order algorithms to solve a soft SAM objective parameterized by a tilting parameter. We provide precise characterizations of the sharpness notions of the tilted SAM framework. Practically, our approach can be used as a gradient-free and memory-efficient alternative to SAM variants, and it achieves better generalization compared to vanilla zeroth-order baselines on a wide range of downstream tasks, including classification, multiple choice QA, and language generation.
Private Zeroth-Order Optimization with Public Data

Xuchen Gong, and Tian Li

NeurIPS 2025

Abstract Paper Code Poster Slides

One of the major bottlenecks for deploying popular first-order differentially private (DP) machine learning algorithms (e.g., DP-SGD) lies in their high computation and memory cost, despite the existence of optimized implementations. Zeroth-order methods have promise in mitigating the overhead, as they leverage function evaluations to approximate the gradients, hence significantly easier to privatize. While recent works have explored zeroth-order approaches in both private and non-private settings, they still suffer from relatively low utilities compared with DP-SGD and limited application domains. In this work, we propose to leverage public information to guide and improve gradient approximation of private zeroth-order algorithms. We explore a suite of \underlinepublic data \underlineassisted \underlinezeroth-\underlineorder optimizers (PAZO) with minimal overhead. We provide theoretical analyses of the PAZO framework under an assumption of the similarity between public and private data. Empirically, we demonstrate that PAZO achieves stronger privacy/utility tradeoffs across vision and text tasks in both pre-training and fine-tuning regimes, outperforming the best first-order baselines (with public gradients) especially in highly private regimes, while offering up to runtime speedup.