Yujia Jin

ICML Conference 2025 Conference Paper

Catoni Contextual Bandits are Robust to Heavy-tailed Rewards

• Chenlu Ye
• Yujia Jin
• Alekh Agarwal
• Tong Zhang 0001

Typical contextual bandit algorithms assume that the rewards at each round lie in some fixed range $[0, R]$, and their regret scales polynomially with this reward range $R$. However, many practical scenarios naturally involve heavy-tailed rewards or rewards where the worst-case range can be substantially larger than the variance. In this paper, we develop an algorithmic approach building on Catoni’s estimator from robust statistics, and apply it to contextual bandits with general function approximation. When the variance of the reward at each round is known, we use a variance-weighted regression approach and establish a regret bound that depends only on the cumulative reward variance and logarithmically on the reward range $R$ as well as the number of rounds $T$. For the unknown-variance case, we further propose a careful peeling-based algorithm and remove the need for cumbersome variance estimation. With additional dependence on the fourth moment, our algorithm also enjoys a variance-based bound with logarithmic reward-range dependence. Moreover, we demonstrate the optimality of the leading-order term in our regret bound through a matching lower bound.

SODA Conference 2024 Conference Paper

A Whole New Ball Game: A Primal Accelerated Method for Matrix Games and Minimizing the Maximum of Smooth Functions

• Yair Carmon
• Arun Jambulapati
• Yujia Jin
• Aaron Sidford

NeurIPS Conference 2024 Conference Paper

Truncated Variance Reduced Value Iteration

• Yujia Jin
• Ishani Karmarkar
• Aaron Sidford
• Jiayi Wang

We provide faster randomized algorithms for computing an $\epsilon$-optimal policy in a discounted Markov decision process with $A_{\text{tot}}$-state-action pairs, bounded rewards, and discount factor $\gamma$. We provide an $\tilde{O}(A_{\text{tot}}[(1 - \gamma)^{-3}\epsilon^{-2} + (1 - \gamma)^{-2}])$-time algorithm in the sampling setting, where the probability transition matrix is unknown but accessible through a generative model which can be queried in $\tilde{O}(1)$-time, and an $\tilde{O}(s + (1-\gamma)^{-2})$-time algorithm in the offline setting where the probability transition matrix is known and $s$-sparse. These results improve upon the prior state-of-the-art which either ran in $\tilde{O}(A_{\text{tot}}[(1 - \gamma)^{-3}\epsilon^{-2} + (1 - \gamma)^{-3}])$ time [Sidford, Wang, Wu, Ye 2018] in the sampling setting, $\tilde{O}(s + A_{\text{tot}} (1-\gamma)^{-3})$ time [Sidford, Wang, Wu, Yang, Ye 2018] in the offline setting, or time at least quadratic in the number of states using interior point methods for linear programming. We achieve our results by building upon prior stochastic variance-reduced value iteration methods [Sidford, Wang, Wu, Yang, Ye 2018]. We provide a variant that carefully truncates the progress of its iterates to improve the variance of new variance-reduced sampling procedures that we introduce to implement the steps. Our method is essentially model-free and can be implemented in $\tilde{O}(A_{\text{tot}})$-space when given generative model access. Consequently, our results take a step in closing the sample-complexity gap between model-free and model-based methods.

ICML Conference 2023 Conference Paper

Quantum Speedups for Zero-Sum Games via Improved Dynamic Gibbs Sampling

• Adam Bouland
• Yosheb M. Getachew
• Yujia Jin
• Aaron Sidford
• Kevin Tian

We give a quantum algorithm for computing an $\epsilon$-approximate Nash equilibrium of a zero-sum game in a $m \times n$ payoff matrix with bounded entries. Given a standard quantum oracle for accessing the payoff matrix our algorithm runs in time $\widetilde{O}(\sqrt{m + n}\cdot \epsilon^{-2. 5} + \epsilon^{-3})$ and outputs a classical representation of the $\epsilon$-approximate Nash equilibrium. This improves upon the best prior quantum runtime of $\widetilde{O}(\sqrt{m + n} \cdot \epsilon^{-3})$ obtained by [van Apeldoorn, Gilyen ’19] and the classical $\widetilde{O}((m + n) \cdot \epsilon^{-2})$ runtime due to [Grigoradis, Khachiyan ’95] whenever $\epsilon = \Omega((m +n)^{-1})$. We obtain this result by designing new quantum data structures for efficiently sampling from a slowly-changing Gibbs distribution.

FOCS Conference 2023 Conference Paper

ReSQueing Parallel and Private Stochastic Convex Optimization

• Yair Carmon
• Arun Jambulapati
• Yujia Jin
• Yin Tat Lee
• Daogao Liu
• Aaron Sidford
• Kevin Tian

We introduce a new tool for stochastic convex optimization (SCO): a Reweighted Stochastic Query (ReSQue) estimator for the gradient of a function convolved with a (Gaussian) probability density. Combining ReSQue with recent advances in ball oracle acceleration [CJJ+20], [ACJ+21], we develop algorithms achieving state-of-the-art complexities for SCO in parallel and private settings. For a SCO objective constrained to the unit ball in $\mathbb{R}^{d}$, we obtain the following results (up to polylogarithmic factors). 1)We give a parallel algorithm obtaining optimization error $\epsilon_{\text {opt} }$ with $d^{1 / 3} \epsilon_{\text {opt} }^{-2 / 3}$ gradient oracle query depth and $d^{1 / 3} \epsilon_{\text {opt} }^{-2 / 3}+\epsilon_{\text {opt} }^{-2}$ gradient queries in total, assuming access to a bounded-variance stochastic gradient estimator. For $\epsilon_{\text {opt} } \in\left[d^{-1}, d^{-1 / 4}\right]$, our algorithm matches the state-of-the-art oracle depth of [BJL+19] while maintaining the optimal total work of stochastic gradient descent. 2)Given n samples of Lipschitz loss functions, prior works [BFTT19], [BFGT20], [AFKT21], [KLL21] established that if $n \gt rsim d \epsilon_{\mathrm{dp}}^{-2}, \left(\epsilon_{\mathrm{dp}}, \delta\right)$-differential privacy is attained at no asymptotic cost to the SCO utility. However, these prior works all required a superlinear number of gradient queries. We close this gap for sufficiently large $n \gt rsim d^{2} \epsilon_{{\bf d p}}^{-3}$, by using ReSQue to design an algorithm with near-linear gradient query complexity in this regime.

NeurIPS Conference 2022 Conference Paper

Optimal and Adaptive Monteiro-Svaiter Acceleration

• Yair Carmon
• Danielle Hausler
• Arun Jambulapati
• Yujia Jin
• Aaron Sidford

We develop a variant of the Monteiro-Svaiter (MS) acceleration framework that removes the need to solve an expensive implicit equation at every iteration. Consequently, for any $p\ge 2$ we improve the complexity of convex optimization with Lipschitz $p$th derivative by a logarithmic factor, matching a lower bound. We also introduce an MS subproblem solver that requires no knowledge of problem parameters, and implement it as either a second- or first-order method by solving linear systems or applying MinRes, respectively. On logistic regression problems our method outperforms previous accelerated second-order methods, but under-performs Newton's method; simply iterating our first-order adaptive subproblem solver is competitive with L-BFGS.

ICML Conference 2022 Conference Paper

RECAPP: Crafting a More Efficient Catalyst for Convex Optimization

• Yair Carmon
• Arun Jambulapati
• Yujia Jin
• Aaron Sidford

The accelerated proximal point method (APPA), also known as "Catalyst", is a well-established reduction from convex optimization to approximate proximal point computation (i. e. , regularized minimization). This reduction is conceptually elegant and yields strong convergence rate guarantees. However, these rates feature an extraneous logarithmic term arising from the need to compute each proximal point to high accuracy. In this work, we propose a novel Relaxed Error Criterion for Accelerated Proximal Point (RECAPP) that eliminates the need for high accuracy subproblem solutions. We apply RECAPP to two canonical problems: finite-sum and max-structured minimization. For finite-sum problems, we match the best known complexity, previously obtained by carefully-designed problem-specific algorithms. For minimizing max_y f(x, y) where f is convex in x and strongly-concave in y, we improve on the best known (Catalyst-based) bound by a logarithmic factor.

SODA Conference 2022 Conference Paper

Semi-Streaming Bipartite Matching in Fewer Passes and Optimal Space

• Sepehr Assadi
• Arun Jambulapati
• Yujia Jin
• Aaron Sidford
• Kevin Tian

We provide Õ ( ∊ –1 )-pass semi-streaming algorithms for computing (1– ∊ )-approximate maximum cardinality matchings in bipartite graphs. Our most efficient methods are deterministic and use optimal, O ( n ), space, improving upon the space complexity of the previous state-of-the-art Õ ( ∊ –1 )-pass algorithm of [AG18]. To obtain our results we provide semi-streaming adaptations of more general continuous optimization tools. Further, we leverage these techniques to obtain improvements for streaming variants of approximate linear programming, optimal transport, exact matching, transshipment, and shortest path problems.

NeurIPS Conference 2021 Conference Paper

Stochastic Bias-Reduced Gradient Methods

• Hilal Asi
• Yair Carmon
• Arun Jambulapati
• Yujia Jin
• Aaron Sidford

We develop a new primitive for stochastic optimization: a low-bias, low-cost estimator of the minimizer $x_\star$ of any Lipschitz strongly-convex function $f$. In particular, we use a multilevel Monte-Carlo approach due to Blanchet and Glynn to turn any optimal stochastic gradient method into an estimator of $x_\star$ with bias $\delta$, variance $O(\log(1/\delta))$, and an expected sampling cost of $O(\log(1/\delta))$ stochastic gradient evaluations. As an immediate consequence, we obtain cheap and nearly unbiased gradient estimators for the Moreau envelope of any Lipschitz convex function. We demonstrate the potential of our estimator through four applications. First, we develop a method for minimizing the maximum of $N$ functions, improving on recent results and matching a lower bound up to logarithmic factors. Second and third, we recover state-of-the-art rates for projection-efficient and gradient-efficient optimization using simple algorithms with a transparent analysis. Finally, we show that an improved version of our estimator would yield a nearly linear-time, optimal-utility, differentially-private non-smooth stochastic optimization method.

ICML Conference 2021 Conference Paper

Towards Tight Bounds on the Sample Complexity of Average-reward MDPs

• Yujia Jin
• Aaron Sidford

We prove new upper and lower bounds for sample complexity of finding an $\epsilon$-optimal policy of an infinite-horizon average-reward Markov decision process (MDP) given access to a generative model. When the mixing time of the probability transition matrix of all policies is at most $t_\mathrm{mix}$, we provide an algorithm that solves the problem using $\widetilde{O}(t_\mathrm{mix} \epsilon^{-3})$ (oblivious) samples per state-action pair. Further, we provide a lower bound showing that a linear dependence on $t_\mathrm{mix}$ is necessary in the worst case for any algorithm which computes oblivious samples. We obtain our results by establishing connections between infinite-horizon average-reward MDPs and discounted MDPs of possible further utility.

NeurIPS Conference 2020 Conference Paper

Acceleration with a Ball Optimization Oracle

• Yair Carmon
• Arun Jambulapati
• Qijia Jiang
• Yujia Jin
• Yin Tat Lee
• Aaron Sidford
• Kevin Tian

Consider an oracle which takes a point x and returns the minimizer of a convex function f in an l 2 ball of radius r around x. It is straightforward to show that roughly r^{-1}\log(1/epsilon) calls to the oracle suffice to find an \epsilon-approximate minimizer of f in an l 2 unit ball. Perhaps surprisingly, this is not optimal: we design an accelerated algorithm which attains an epsilon-approximate minimizer with roughly r^{-2/3} \log(1/epsilon) oracle queries, and give a matching lower bound. Further, we implement ball optimization oracles for functions with a locally stable Hessian using a variant of Newton's method and, in certain cases, stochastic first-order methods. The resulting algorithms apply to a number of problems of practical and theoretical import, improving upon previous results for logistic and l infinity regression and achieving guarantees comparable to the state-of-the-art for l p regression.

FOCS Conference 2020 Conference Paper

Coordinate Methods for Matrix Games

• Yair Carmon
• Yujia Jin
• Aaron Sidford
• Kevin Tian

We develop primal-dual coordinate methods for solving bilinear saddle-point problems of the form minx ∈ Xmaxy ∈ Yy T Ax which contain linear programming, classification, and regression as special cases. Our methods push existing fully stochastic sublinear methods and variance-reduced methods towards their limits in terms of per-iteration complexity and sample complexity. We obtain nearly-constant per-iteration complexity by designing efficient data structures leveraging Taylor approximations to the exponential and a binomial heap. We improve sample complexity via low-variance gradient estimators using dynamic sampling distributions that depend on both the iterates and the magnitude of the matrix entries. Our runtime bounds improve upon those of existing primal-dual methods by a factor depending on sparsity measures of the m by n matrix A. For example, when rows and columns have constant l1/l2 norm ratios, we offer improvements by a factor of m+n in the fully stochastic setting and √{m+n} in the variance-reduced setting. We apply our methods to computational geometry problems, i. e. minimum enclosing ball, maximum inscribed ball, and linear regression, and obtain improved complexity bounds. For linear regression with an elementwise nonnegative matrix, our guarantees improve on exact gradient methods by a factor of √{nnz(A)/(m+n)}.

ICML Conference 2020 Conference Paper

Efficiently Solving MDPs with Stochastic Mirror Descent

• Yujia Jin
• Aaron Sidford

We present a unified framework based on primal-dual stochastic mirror descent for approximately solving infinite-horizon Markov decision processes (MDPs) given a generative model. When applied to an average-reward MDP with $A_{tot}$ total actions and mixing time bound $t_{mix}$ our method computes an $\epsilon$-optimal policy with an expected $\widetilde{O}(t_{mix}^2 A_{tot} \epsilon^{-2})$ samples from the state-transition matrix, removing the ergodicity dependence of prior art. When applied to a $\gamma$-discounted MDP with $A_{tot}$ total actions our method computes an $\epsilon$-optimal policy with an expected $\widetilde{O}((1-\gamma)^{-4} A_{tot} \epsilon^{-2})$ samples, improving over the best-known primal-dual methods while matching the state-of-the-art up to a $(1-\gamma)^{-1}$ factor. Both methods are model-free, update state values and policies simultaneously, and run in time linear in the number of samples taken. We achieve these results through a more general stochastic mirror descent framework for solving bilinear saddle-point problems with simplex and box domains and we demonstrate the flexibility of this framework by providing further applications to constrained MDPs.

NeurIPS Conference 2019 Conference Paper

Principal Component Projection and Regression in Nearly Linear Time through Asymmetric SVRG

• Yujia Jin
• Aaron Sidford

Given a n-by-d data matrix A, principal component projection (PCP) and principal component regression (PCR), i. e. projection and regression restricted to the top-eigenspace of A, are fundamental problems in machine learning, optimization, and numerical analysis. In this paper we provide the first algorithms that solve these problems in nearly linear time for fixed eigenvalue distribution and large n. This improves upon previous methods which had superlinear running times when either the number of top eigenvalues or gap between the eigenspaces were large. We achieve our results by applying rational polynomial approximations to reduce the problem to solving asymmetric linear systems which we solve by a variant of SVRG. We corroborate these findings with preliminary empirical experiments.

NeurIPS Conference 2019 Conference Paper

Variance Reduction for Matrix Games

• Yair Carmon
• Yujia Jin
• Aaron Sidford
• Kevin Tian

We present a randomized primal-dual algorithm that solves the problem min x max y y^T A x to additive error epsilon in time nnz(A) + sqrt{nnz(A) n} / epsilon, for matrix A with larger dimension n and nnz(A) nonzero entries. This improves the best known exact gradient methods by a factor of sqrt{nnz(A) / n} and is faster than fully stochastic gradient methods in the accurate and/or sparse regime epsilon < sqrt{n / nnz(A)$. Our results hold for x, y in the simplex (matrix games, linear programming) and for x in an \ell_2 ball and y in the simplex (perceptron / SVM, minimum enclosing ball). Our algorithm combines the Nemirovski's "conceptual prox-method" and a novel reduced-variance gradient estimator based on "sampling from the difference" between the current iterate and a reference point.

TCS Journal 2017 Journal Article

Maximum matchings and minimum dominating sets in Apollonian networks and extended Tower of Hanoi graphs

• Yujia Jin
• Huan Li
• Zhongzhi Zhang