
I am a research scientist at Meta. I obtained my PhD from the Department of Computer Sciences, UW-Madison. I was fortunate to be advised by Prof. Jelena Diakonikolas.
Earlier, I received my BS in Math from Zhiyuan Honors College at Shanghai Jiao Tong University.
My research interests are primarily in optimization and machine learning.
Please reach out via xcai74 [at] wisc [dot] edu.
Recent Papers
(reverse chronological, * denotes equal contribution)
OneShot: Index-in-Ranking with Neural Scoring for Large-Scale Retrieval.
Ziwei Li, Shuyao Li, Xufeng Cai,
Xue Zou, Yiming Ma,
Huiting Lu, Wujie Yan, Zhichen Zhao, Yang Lu, Zhe Wang,
Rui Luo, Zhengyu Su,
Dan Zhang, Yimin Tan, Ji Liu.
arXiv preprint, arXiv:2607.27475, 2026.
abstract / arXiv
In modern recommendation systems, retrieval serves as a primary stage responsible for filtering billions of candidate items down to thousands prior to refined ranking.
To make this massive search effective and efficient, the system relies on ranking accuracy and indexing efficiency. However, these two objectives are traditionally misaligned:
while the former optimizes for the alignment between ranking predictions and user behavior, the latter optimizes for a structural grouping of item representations which enables fast search among billions of candidates.
Thus, despite extensive efforts to scale up interaction modeling for retrieval, they remain fundamentally limited by the structural misalignment between the ranking objectives and the proximity-learned index.
In this work, we address this long-standing dichotomy by proposing a new holistic retrieval framework, OneShot. It is an end-to-end, in-model index learning framework that natively aligns index learning with ranking objectives.
Using this joint learning as a structural foundation, OneShot pushes the boundaries of retrieval expressiveness by scaling interaction modeling with neural scoring beyond the persistent dot-product bottleneck.
OneShot is fully deployed in Instagram's industrial short-video recommendation system, driving significant wins in user daily sessions, engagement, and time-spent.
Additionally, OneShot achieves a 20% recall gain at the operational ranking volume and a 10x efficiency improvement at an equivalent recall level.
Adaptive Delayed-Update Cyclic Algorithm for Variational Inequalities.
Yi Wei, Xufeng Cai, Jelena Diakonikolas.
In Proc. NeurIPS'26, 2026.
abstract / arXiv
Cyclic block coordinate methods are a fundamental class of first-order algorithms, widely used in practice for their simplicity and strong empirical performance.
Yet, their theoretical behavior remains challenging to explain, and setting their step sizes -- beyond classical coordinate descent for minimization -- typically requires careful tuning or line-search machinery.
In this work, we develop ADUCA (Adaptive Delayed-Update Cyclic Algorithm), a cyclic algorithm addressing a broad class of Minty variational inequalities with monotone Lipschitz operators.
ADUCA is parameter-free: it requires no global or block-wise Lipschitz constants and uses no per-epoch line search, except at initialization.
A key feature of the algorithm is using operator information delayed by a full cycle, which makes the algorithm compatible with parallel and distributed implementations, and attractive due to weakened synchronization requirements across blocks.
We prove that ADUCA attains (near) optimal global oracle complexity as a function of target error \(\epsilon\)>0, scaling with \(1/\epsilon\) for monotone operators, or with \(\log^2(1/\epsilon)\) for operators that are strongly monotone.
Near-Linear Runtime for a Classical Matrix Preconditioning Algorithm.
Xufeng Cai, Jason M. Altschuler, Jelena Diakonikolas.
Journal of the ACM (JACM), 73(1), pp. 1-25, 2026.
abstract / arXiv
In 1960, Osborne proposed a simple iterative algorithm for matrix balancing with outstanding numerical performance.
Today, it is the default preconditioning procedure before eigenvalue computation and other linear algebra subroutines in mainstream software packages such as Python, Julia, MATLAB, EISPACK, LAPACK, and more.
Despite its widespread usage, Osborne's algorithm has long resisted theoretical guarantees for its runtime: the first polynomial-time guarantees were obtained only in the past decade,
and recent near-linear runtimes remain confined to variants of Osborne's algorithm with important differences that make them simpler to analyze but empirically slower.
In this paper, we address this longstanding gap between theory and practice by proving that Osborne's original algorithm -- the de facto preconditioner in practice -- in fact has a near-linear runtime.
This runtime guarantee (1) is optimal in the input size up to at most a single logarithm, (2) is the first runtime for Osborne's algorithm that does not dominate the runtime of downstream tasks like eigenvalue computation,
and (3) improves upon the theoretical runtimes for all other variants of Osborne's algorithm.
Last Iterate Convergence of Incremental Methods as a Model of Forgetting.
Xufeng Cai, Jelena Diakonikolas.
In Proc. ICLR'25, 2025.
abstract / arXiv
Incremental gradient and incremental proximal methods are a fundamental class of optimization algorithms used for solving finite sum problems,
broadly studied in the literature. Yet, without strong convexity, their convergence guarantees have primarily been established for the ergodic (average) iterate.
We establish the first nonasymptotic convergence guarantees for the last iterate of both incremental gradient and incremental proximal methods,
in general convex smooth (for both) and convex Lipschitz (for the proximal variants) settings. Our oracle complexity bounds for the last iterate nearly match
(i.e., match up to a square-root-log or a log factor) the best known oracle complexity bounds for the average iterate, for both classes of methods.
We further obtain generalizations of our results to weighted averaging of the iterates with increasing weights and for randomly permuted ordering of updates.
We study last iterate convergence of the incremental proximal method as a mathematical abstraction of forgetting in continual learning and prove a lower bound
that certifies that a large amount of regularization is crucial to mitigating catastrophic forgetting---one of the key considerations in continual learning.
Our results generalize last iterate guarantees for incremental methods compared to state of the art, as such results were previously known only for overparameterized linear models,
which correspond to convex quadratic problems with infinitely many solutions.
Tighter Convergence Bounds for Shuffled SGD via Primal-Dual Perspective.
Xufeng Cai*, Cheuk Yin Lin*, Jelena Diakonikolas.
In Proc. NeurIPS'24, 2024.
abstract / arXiv
Stochastic gradient descent (SGD) is perhaps the most prevalent optimization method in modern
machine learning. Contrary to the empirical practice of sampling from the datasets without replacement and
with (possible) reshuffling at each epoch, the theoretical counterpart of SGD usually relies on the assumption
of sampling with replacement. It is only very recently that SGD with sampling without replacement – shuffled SGD – has been analyzed. For convex finite sum problems with n components and under the
\(L\)-smoothness assumption for each component function, there are matching upper and lower bounds, under
sufficiently small – \(\mathcal{O}(\frac{1}{nL})\) – step sizes. Yet those bounds appear too pessimistic – in fact, the predicted
performance is generally no better than for full gradient descent – and do not agree with the empirical
observations. In this work, to narrow the gap between the theory and practice of shuffled SGD, we sharpen
the focus from general finite sum problems to empirical risk minimization with linear predictors. This
allows us to take a primal-dual perspective and interpret shuffled SGD as a primal-dual method with cyclic
coordinate updates on the dual side. Leveraging this perspective, we prove a fine-grained complexity bound
that depends on the data matrix and is never worse than what is predicted by the existing bounds. Notably,
our bound can predict much faster convergence than the existing analyses – by a factor of the order of \(\sqrt{n}\)
in some cases. We empirically demonstrate that on common machine learning datasets our bound is indeed
much tighter. We further show how to extend our analysis to convex nonsmooth problems, with similar improvements.
Variance Reduced Halpern Iteration for Finite-Sum Monotone Inclusions.
Xufeng Cai*, Ahmet Alacaoglu*, Jelena Diakonikolas.
In Proc. ICLR'24, 2024.
abstract / arXiv
Machine learning approaches relying on such criteria as adversarial robustness or multi-agent settings have raised the need for solving game-theoretic equilibrium problems.
Of particular relevance to these applications are methods targeting finite-sum structure, which generically arises in empirical variants of learning problems in these contexts.
Further, methods with computable approximation errors are highly desirable, as they provide verifiable exit criteria. Motivated by these applications,
we study finite-sum monotone inclusion problems, which model broad classes of equilibrium problems.
Our main contributions are variants of the classical Halpern iteration that employ variance reduction to obtain improved complexity guarantees in which \(n\) component operators in the finite sum are ``on average''
either cocoercive or Lipschitz continuous and monotone, with parameter \(L\). The resulting oracle complexity of our methods, which provide guarantees for the last iterate and for a (computable) operator norm residual,
is \(\widetilde{\mathcal{O}}( n + \sqrt{n}L\varepsilon^{-1})\), which improves upon existing methods by a factor up to \(\sqrt{n}\).
This constitutes the first variance reduction-type result for general finite-sum monotone inclusions and for more specific problems such as convex-concave optimization when operator norm residual is the optimality measure.
We further argue that, up to poly-logarithmic factors, this complexity is unimprovable in the monotone Lipschitz setting; i.e., the provided result is near-optimal.
Cyclic Block Coordinate Descent With Variance Reduction for Composite Nonconvex Optimization.
Xufeng Cai, Chaobing Song, Stephen J. Wright, Jelena Diakonikolas.
In Proc. ICML'23, 2023.
abstract / arXiv
Nonconvex optimization is central in solving many machine learning problems, in which block-wise structure is
commonly encountered. In this work, we propose cyclic block coordinate methods for nonconvex optimization problems
with non-asymptotic gradient norm guarantees. Our convergence analysis is based on a gradient Lipschitz condition
with respect to a Mahalanobis norm, inspired by a recent progress on cyclic block coordinate methods. In deterministic
settings, our convergence guarantee matches the guarantee of (full-gradient) gradient descent, but with the gradient
Lipschitz constant being defined w.r.t. a Mahalanobis norm. In stochastic settings, we use recursive variance reduction to
decrease the per-iteration cost and match the arithmetic operation complexity of current optimal stochastic full-gradient
methods, with a unified analysis for both finite-sum and infinite-sum cases. We prove a faster linear convergence result
when a Polyak-Łojasiewicz (PŁ) condition holds. To our knowledge, this work is the first to provide non-asymptotic
convergence guarantees — variance-reduced or not — for a cyclic block coordinate method in general composite (smooth
+ nonsmooth) nonconvex settings. Our experimental results demonstrate the efficacy of the proposed cyclic scheme in
training deep neural nets.
Stochastic Halpern Iteration with Variance Reduction for Stochastic Monotone Inclusions.
Xufeng Cai, Chaobing Song, Cristóbal Guzmán, Jelena Diakonikolas.
In Proc. NeurIPS'22, 2022.
abstract / arXiv
We study stochastic monotone inclusion problems, which widely appear in machine learning applications,
including robust regression and adversarial learning. We propose novel variants of stochastic Halpern iteration with recursive variance reduction.
In the cocoercive---and more generally Lipschitz-monotone---setup,
our algorithm attains \(\epsilon\) norm of the operator with \(\mathcal{O}(\frac{1}{\epsilon^3})\) stochastic operator evaluations,
which significantly improves over state of the art \(\mathcal{O}(\frac{1}{\epsilon^4})\) stochastic operator evaluations
required for existing monotone inclusion solvers applied to the same problem classes. %is the first result of this kind.
We further show how to couple one of the proposed variants of stochastic Halpern iteration
with a scheduled restart scheme to solve stochastic monotone inclusion problems with \({\mathcal{O}}(\frac{\log(1/\epsilon)}{\epsilon^2})\)
stochastic operator evaluations under additional sharpness or strong monotonicity assumptions.
Recent Talks
Last Iterate Convergence of Incremental Methods as a Model of Forgetting.
International Conference on Continuous Optimization (ICCOPT)
Los Angeles, CA, USA, July 2025.
Variance Reduced Halpern Iteration for Finite-Sum Monotone Inclusions.
INFORMS Annual Meeting (INFORMS)
Seattle, WA, USA, Oct 2024.
Stochastic Halpern Iteration with Variance Reduction for Stochastic Monotone Inclusions.
International Conference on Continuous Optimization (ICCOPT)
Bethlehem, PA, USA, July 2022.
Teaching
CS639, Foundations of Data Science, UW-Madison, Spring 2024 (TA).
CS639, Foundations of Data Science, UW-Madison, Spring 2022 (TA).
CS760, Machine Learning, UW-Madison, Spring 2021 (TA).
CS760, Machine Learning, UW-Madison, Fall 2020 (TA).
Miscellaneous
In my free time, I enjoy reading, photography, and music (vinyls).