EDBT 2026 Demo / reviewers in the wild / expert
Yuejie Chi
dblp:82/8759
· DBLP profile ↗
108ranked-venue papers
11as first author
64since 2021 · last 2025
0000-0002-6766-5459ORCID · verified
Domains — the database's venue-derived domains; a paper can count in several
Artificial intelligence and machine learning · 63 · 2 first-author · 52 since 2021Graphics, computer vision, multimedia, augmented reality and games · 30 · 8 first-author · 9 since 2021Theory of computation · 8 · 4 since 2021Applied, interdisciplinary, general and emerging computing · 6 · 1 first-authorComputer networks · 3 · 1 first-author · 1 since 2021Databases, data management, data science and information retrieval · 2 · 2 since 2021
| Year | Publication | Venue | Position |
|---|---|---|---|
| 2025 | Characterizing the Accuracy-Communication-Privacy Trade-off in Distributed Stochastic Convex OptimizationabstractWe consider the problem of differentially private stochastic convex optimization (DP-SCO) in a distributed setting with $M$ clients, where each of them has a local dataset of $N$ i.i.d. data samples from an underlying data distribution. The objective is to design an algorithm to minimize a convex population loss using a collaborative effort across $M$ clients, while ensuring the privacy of the local datasets. In this work, we investigate the accuracy-communication-privacy trade-off for this problem. We establish matching converse and achievability results using a novel lower bound and a new algorithm for distributed DP-SCO based on Vaidya’s plane cutting method. Thus, our results provide a complete characterization of the accuracy-communication-privacy trade-off for DP-SCO in the distributed setting. Sudeep Salgia, Nikola Pavlovic, Yuejie Chi, Qing Zhao 0001 |
AISTATS | 3 |
| 2025 | Faster WIND: Accelerating Iterative Best-of-N Distillation for LLM AlignmentabstractRecent advances in aligning large language models with human preferences have corroborated the growing importance of best-of-$N$ distillation (BOND). However, the iterative BOND algorithm is prohibitively expensive in practice due to the sample and computation inefficiency. This paper addresses the problem by revealing a unified game-theoretic connection between iterative BOND and self-play alignment, which unifies seemingly disparate algorithmic paradigms. Based on the connection, we establish a novel framework, \textbf{WIN} rate \textbf{D}ominance (WIND), with a series of efficient algorithms for regularized win rate dominance optimization that approximates iterative BOND in the parameter space. We provides provable sample efficiency guarantee for one of the WIND variant with the square loss objective. The experimental results confirm that our algorithm not only accelerates the computation, but also achieves superior sample efficiency compared to existing methods. Tong Yang 0007, Jincheng Mei, Hanjun Dai, Zixin Wen, Shicong Cen, Dale Schuurmans, Yuejie Chi, Bo Dai 0001 |
AISTATS | 7 |
| 2025 | Leveraging Multimodal Diffusion Models to Accelerate Imaging with Side InformationabstractDiffusion models have found phenomenal success as expressive priors for solving inverse problems, but their extension beyond natural images to more structured scientific domains remains limited. Motivated by applications in materials science, we aim to reduce the number of measurements required from an expensive imaging modality of interest, by leveraging side information from an auxiliary modality that is much cheaper to obtain. To deal with the non-differentiable and black-box nature of the forward model, we propose a framework to train a multimodal diffusion model over the joint modalities, turning inverse problems with black-box forward models into simple linear inpainting problems. Numerically, we demonstrate the feasibility of training diffusion models over materials imagery data, and show that our approach achieves superior image reconstruction by leveraging the available side information, requiring significantly less amount of data from the expensive microscopy modality. Timofey Efimov, Harry Dong, Megna Shah, Jeff P. Simmons, Sean Donegan, Yuejie Chi |
ICASSP | 6 |
| 2025 | A Theoretical Analysis of Self-Supervised Learning for Vision TransformersabstractSelf-supervised learning has become a cornerstone in computer vision, primarily divided into reconstruction-based methods like masked autoencoders (MAE) and discriminative methods such as contrastive learning (CL). Recent empirical observations reveal that MAE and CL capture different types of representations: CL tends to focus on global patterns, while MAE adeptly captures **both global and subtle local** information simultaneously. Despite a flurry of recent empirical investigations to shed light on this difference, theoretical understanding remains limited, especially on the dominant architecture **vision transformers** (ViTs). In this paper, to provide rigorous insights, we model the visual data distribution by considering two types of spatial features: dominant global features and comparatively minuscule local features, and study the impact of imbalance among these features. We analyze the training dynamics of one-layer softmax-based ViTs on both MAE and CL objectives using gradient descent. Our analysis shows that as the degree of feature imbalance varies, ViTs trained with the MAE objective effectively learn both global and local features to achieve near-optimal reconstruction, while the CL-trained ViTs favor predominantly global features, even under mild imbalance. These results provide a theoretical explanation for distinct behaviors of MAE and CL observed in empirical studies. Yu Huang 0023, Zixin Wen, Yuejie Chi, Yingbin Liang |
ICLR | 3 |
| 2025 | Value-Incentivized Preference Optimization: A Unified Approach to Online and Offline RLHFabstractReinforcement learning from human feedback (RLHF) has demonstrated great promise in aligning large language models (LLMs) with human preference. Depending on the availability of preference data, both online and offline RLHF are active areas of investigation. A key bottleneck is understanding how to incorporate uncertainty estimation in the reward function learned from the preference data for RLHF, regardless of how the preference data is collected. While the principles of optimism or pessimism under uncertainty are well-established in standard reinforcement learning (RL), a practically-implementable and theoretically-grounded form amenable to large language models is not yet available, as standard techniques for constructing confidence intervals become intractable under arbitrary policy parameterizations.
In this paper, we introduce a unified approach to online and offline RLHF --- value-incentivized preference optimization (VPO) --- which regularizes the maximum-likelihood estimate of the reward function with the corresponding value function, modulated by a sign to indicate whether the optimism or pessimism is chosen. VPO also directly optimizes the policy with implicit reward modeling, and therefore shares a simpler RLHF pipeline similar to direct preference optimization. Theoretical guarantees of VPO are provided for both online and offline settings, matching the rates of their standard RL counterparts. Moreover, experiments on text summarization, dialogue, and standard benchmarks verify the practicality and effectiveness of VPO. Shicong Cen, Jincheng Mei, Katayoon Goshvadi, Hanjun Dai, Tong Yang 0007, Sherry Yang 0001, Dale Schuurmans, Yuejie Chi, Bo Dai 0001 |
ICLR | 8 |
| 2025 | Robust Gymnasium: A Unified Modular Benchmark for Robust Reinforcement LearningabstractDriven by inherent uncertainty and the sim-to-real gap, robust reinforcement learning (RL) seeks to improve resilience against the complexity and variability in agent-environment sequential interactions. Despite the existence of a large number of RL benchmarks, there is a lack of standardized benchmarks for robust RL. Current robust RL policies often focus on a specific type of uncertainty and are evaluated in distinct, one-off environments. In this work, we introduce Robust-Gymnasium, a unified modular benchmark designed for robust RL that supports a wide variety of disruptions across all key RL components—agents' observed state and reward, agents' actions, and the environment. Offering over sixty diverse task environments spanning control and robotics, safe RL, and multi-agent RL, it provides an open-source and user-friendly tool for the community to assess current methods and foster the development of robust RL algorithms.
In addition, we benchmark existing standard and robust RL algorithms within this framework, uncovering significant deficiencies in each and offering new insights. Shangding Gu, Laixi Shi, Muning Wen, Ming Jin 0002, Eric Mazumdar, Yuejie Chi, Adam Wierman, Costas J. Spanos |
ICLR | 6 |
| 2025 | Vertical Federated Learning with Missing Features During Training and InferenceabstractVertical federated learning trains models from feature-partitioned datasets across multiple clients, who collaborate without sharing their local data. Standard approaches assume that all feature partitions are available during both training and inference. Yet, in practice, this assumption rarely holds, as for many samples only a subset of the clients observe their partition. However, not utilizing incomplete samples during training harms generalization, and not supporting them during inference limits the utility of the model. Moreover, if any client leaves the federation after training, its partition becomes unavailable, rendering the learned model unusable. Missing feature blocks are therefore a key challenge limiting the applicability of vertical federated learning in real-world scenarios. To address this, we propose LASER-VFL, a vertical federated learning method for efficient training and inference of split neural network-based models that is capable of handling arbitrary sets of partitions. Our approach is simple yet effective, relying on the sharing of model parameters and on task-sampling to train a family of predictors. We show that LASER-VFL achieves a $\mathcal{O}({1}/{\sqrt{T}})$ convergence rate for nonconvex objectives and, under the Polyak-Łojasiewicz inequality, it achieves linear convergence to a neighborhood of the optimum. Numerical experiments show improved performance of LASER-VFL over the baselines. Remarkably, this is the case even in the absence of missing features. For example, for CIFAR-100, we see an improvement in accuracy of $19.3$\% when each of four feature blocks is observed with a probability of 0.5 and of $9.5$\% when all features are observed. The code for this work is available at https://github.com/Valdeira/LASER-VFL. Pedro Valdeira, Yuejie Chi |
ICLR | 3 |
| 2025 | Incentivize without Bonus: Provably Efficient Model-based Online Multi-agent RL for Markov GamesabstractMulti-agent reinforcement learning (MARL) lies at the heart of a plethora of applications involving the interaction of a group of agents in a shared unknown environment. A prominent framework for studying MARL is Markov games, with the goal of finding various notions of equilibria in a sample-efficient manner, such as the Nash equilibrium (NE) and the coarse correlated equilibrium (CCE). However, existing sample-efficient approaches either require tailored uncertainty estimation under function approximation, or careful coordination of the players. In this paper, we propose a novel model-based algorithm, called VMG, that incentivizes exploration via biasing the empirical
estimate of the model parameters towards those with a higher collective best-response values of all the players when fixing the other players' policies, thus encouraging the policy to deviate from its current equilibrium for more exploration. VMG is oblivious to different forms of function approximation, and permits simultaneous and uncoupled policy updates of all players. Theoretically, we also establish that VMG achieves a near-optimal regret for finding both the NEs of two-player zero-sum Markov games and CCEs of multi-player general-sum Markov games under linear function approximation in an online environment, which nearly match their counterparts with sophisticated uncertainty quantification. Tong Yang 0007, Bo Dai 0001, Yuejie Chi |
ICML | 4 |
| 2025 | Breaking the Curse of Multiagency in Robust Multi-Agent Reinforcement LearningabstractStandard multi-agent reinforcement learning (MARL) algorithms are vulnerable to sim-to-real gaps. To address this, distributionally robust Markov games (RMGs) have been proposed to enhance robustness in MARL by optimizing the worst-case performance when game dynamics shift within a prescribed uncertainty set. RMGs remains under-explored, from reasonable problem formulation to the development of sample-efficient algorithms. Two notorious and open challenges are the formulation of the uncertainty set and whether the corresponding RMGs can overcome the curse of multiagency, where the sample complexity scales exponentially with the number of agents. In this work, we propose a natural class of RMGs inspired by behavioral economics, where each agent's uncertainty set is shaped by both the environment and the integrated behavior of other agents. We first establish the well-posedness of this class of RMGs by proving the existence of game-theoretic solutions such as robust Nash equilibria and coarse correlated equilibria (CCE). Assuming access to a generative model, we then introduce a sample-efficient algorithm for learning the CCE whose sample complexity scales polynomially with all relevant parameters. To the best of our knowledge, this is the first algorithm to break the curse of multiagency for RMGs, regardless of the uncertainty set formulation. Laixi Shi, Jingchu Gai, Eric Mazumdar, Yuejie Chi, Adam Wierman |
ICML | 4 |
| 2025 | ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM InferenceabstractWith the widespread deployment of long-context large language models (LLMs), there has been a growing demand for efficient support of high-throughput inference. However, as the key-value (KV) cache expands with the sequence length, the increasing memory footprint and the need to access it for decoding both result in low throughput when serving long-context LLMs. While various dynamic sparse attention methods have been proposed to accelerate inference while maintaining generation quality, they either fail to sufficiently reduce GPU memory usage or introduce significant decoding latency by offloading the KV cache to the CPU. We present ShadowKV, a high-throughput long-context LLM inference system that stores the low-rank key cache and offloads the value cache to reduce the memory footprint for larger batch sizes and longer sequences. To minimize decoding latency, ShadowKV employs an accurate KV selection strategy that reconstructs minimal sparse KV pairs on-the-fly. By evaluating ShadowKV on benchmarks like RULER, LongBench, and models such as Llama-3.1-8B and GLM-4-9B-1M, we demonstrate that it achieves up to 6$\times$ larger batch sizes and 3.04$\times$ higher throughput on an A100 GPU without sacrificing accuracy, even surpassing the performance achievable with infinite batch size under the assumption of infinite GPU memory. Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng 0001, Ningxin Zheng, Xin Liu 0086, Harry Dong, Yuejie Chi, Beidi Chen |
ICML | 8 |
| 2025 | Transformers Provably Learn Chain-of-Thought Reasoning with Length GeneralizationabstractThe ability to reason lies at the core of artificial intelligence (AI), and challenging problems usually call for deeper and longer reasoning to tackle. A crucial question about AI reasoning is whether models can extrapolate learned reasoning patterns to solve harder tasks with a longer chain-of-thought (CoT). In this work, we present a theoretical analysis of transformers learning on synthetic state-tracking tasks with gradient descent. Specifically: 1). We prove how the *algebraic structure* of state-tracking problems governs the length generalization of learned reasoning in transformers. In doing so, we formulate the **attention concentration** mechanism, linking the retrieval robustness of the attention layer to the task structure of long-context state tracking problems. 2). Moreover, we prove that a transformer can provably *self-improve* via a *recursive self-training* scheme that progressively extends the range of solvable problem lengths. We show that the model can achieve abilities outside the coverage of the base model in recursive training, different from prior theoretical works on self-improvement.
To our knowledge, we provide the first *optimization guarantee* that constant-depth transformers provably learn $\text{NC}^1$-complete problems with CoT, significantly going beyond prior art confined in $\text{TC}^0$, unless the widely held conjecture $\text{TC}^0 \neq \text{NC}^1$ fails. Finally, we present a broad set of experiments supporting our theoretical results, confirming the length generalization behaviors and the mechanism of attention concentration. Yu Huang 0023, Zixin Wen, Aarti Singh, Yuejie Chi, Yuxin Chen 0002 |
NeurIPS | 4 |
| 2025 | Exploration from a Primal-Dual Lens: Value-Incentivized Actor-Critic Methods for Sample-Efficient Online RLabstractOnline reinforcement learning (RL) with complex function approximations such as transformers and deep neural networks plays a significant role in the modern practice of artificial intelligence. Despite its popularity and importance, balancing the fundamental trade-off between exploration and exploitation remains a long-standing challenge; in particular, we are still in lack of efficient and practical schemes that are backed by theoretical performance guarantees. Motivated by recent developments in exploration via optimistic regularization, this paper provides an interpretation of the principle of optimism through the lens of primal-dual optimization. From this fresh perspective, we set forth a new value-incentivized actor-critic (VAC) method, which optimizes a single easy-to-optimize objective integrating exploration and exploitation --- it promotes state-action and policy estimates that are both consistent with collected data transitions and result in higher value functions. Theoretically, the proposed VAC method has near-optimal regret guarantees under linear Markov decision processes (MDPs) in both finite-horizon and infinite-horizon settings, which can be extended to the general function approximation setting under appropriate assumptions. Tong Yang 0007, Bo Dai 0001, Yuejie Chi |
NeurIPS | 4 |
| 2025 | Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient DescentabstractTransformers have demonstrated remarkable capabilities in multi-step reasoning tasks. However, understandings of the underlying mechanisms by which they acquire these abilities through training remain limited, particularly from a theoretical standpoint. This work investigates how transformers learn to solve symbolic multi-step reasoning problems through chain-of-thought processes, focusing on path-finding in trees. We analyze two intertwined tasks: a backward reasoning task, where the model outputs a path from a goal node to the root, and a more complex forward reasoning task, where the model implements two-stage reasoning by first identifying the goal-to-root path and then reversing it to produce the root-to-goal path. Our theoretical analysis, grounded in the dynamics of gradient descent, shows that trained one-layer transformers can provably solve both tasks with generalization guarantees to unseen trees. In particular, our multi-phase training dynamics for forward reasoning elucidate how different attention heads learn to specialize and coordinate autonomously to solve the two subtasks in a single autoregressive path. These results provide a mechanistic explanation of how trained transformers can implement sequential algorithmic procedures. Moreover, they offer insights into the emergence of reasoning abilities, suggesting that when tasks are structured to take intermediate chain-of-thought steps, even shallow multi-head transformers can effectively solve problems that would otherwise require deeper architectures. Tong Yang 0007, Yu Huang 0023, Yingbin Liang, Yuejie Chi |
NeurIPS | 4 |
| 2025 | The Blessing of Heterogeneity in Federated Q-Learning: Linear Speedup and BeyondabstractIn this paper, we consider federated Q-learning, which aims to learn an optimal Q-function by periodically aggregating local Q-estimates trained on local data alone. Focusing on infinite-horizon tabular Markov decision processes, we provide sample complexity guarantees for both the synchronous and asynchronous variants of federated Q-learning, which exhibit a linear speedup with respect to the number of agents and near-optimal dependencies on other salient problem parameters. In the asynchronous setting, existing analyses of federated Q-learning, which adopt an equally weighted averaging of local Q-estimates, require that every agent covers the entire state-action space. In contrast, our improved sample complexity scales inverse proportionally to the minimum entry of the average stationary state-action occupancy distribution of all agents, thus only requiring the agents to collectively cover the entire state-action space, unveiling the blessing of heterogeneity. However, its sample complexity still suffers when the local trajectories are highly heterogeneous. In response, we propose a novel federated Q-learning algorithm with importance averaging, giving larger weights to more frequently visited state-action pairs, which achieves a robust linear speedup as if all trajectories are centrally processed, regardless of the heterogeneity of local behavior policies. Jiin Woo, Gauri Joshi, Yuejie Chi |
J. Mach. Learn. Res. | 3 |
| 2024 | Escaping Saddle Points in Heterogeneous Federated Learning via Distributed SGD with Communication CompressionabstractWe consider the problem of finding second-order stationary points in the optimization of heterogeneous federated learning (FL). Previous works in FL mostly focus on first-order convergence guarantees, which do not rule out the scenario of unstable saddle points. Meanwhile, it is a key bottleneck of FL to achieve communication efficiency without compensating the learning accuracy, especially when local data are highly heterogeneous across different clients. Given this, we propose a novel algorithm PowerEF-SGD that only communicates compressed information via a novel error-feedback scheme. To our knowledge, PowerEF-SGD is the first distributed and compressed SGD algorithm that provably escapes saddle points in heterogeneous FL without any data homogeneity assumptions. In particular, PowerEF-SGD improves to second-order stationary points after visiting first-order (possibly saddle) points, using additional gradient queries and communication rounds only of almost the same order required by first-order convergence, and the convergence rate shows a linear-speedup pattern in terms of the number of workers. Our theory improves/recovers previous results, while extending to much more tolerant settings on the local data. Numerical experiments are provided to complement the theory. Sijin Chen, Zhize Li 0001, Yuejie Chi |
AISTATS | 3 |
| 2024 | Scalable Dynamic Resource Allocation via Domain Randomized Reinforcement LearningabstractIn 5G wireless networks, the User Plane Function (UPF) plays a crucial role in efficiently transferring users’ traffic — a series of data packets — to manage internet communications. Setting the server’s processor frequency excessively high can easily meet the packet drop requirements but may lead to unnecessary power consumption. Therefore, as user traffic fluctuates, selecting the optimal processor frequency is essential for minimizing power consumption while satisfying packet drop constraints. This challenge motivates us to address the dynamic resource (frequency) allocation problem, where deep reinforcement learning (RL) has shown significant potential. Most existing studies train and evaluate the RL model in the same environment with consistent traffic patterns. However, frequent variations in user traffic can cause the policy trained on the outdated traffic to fail catastrophically on unseen traffic.To address such traffic distribution shifts, we propose a two-phase RL approach augmented with Automatic Domain Randomization (RL-ADR). This method includes a training phase that utilizes domain randomization to create a library of policy candidates, and an inference phase that selects the optimal frequency using this policy library alongside a safe data buffer. The proposed RL-ADR achieves zero packet drops on two unseen long-horizon traffics (3 hours) after being trained on 25 synthetic traffics that only span for 18 seconds. Compared to static resource allocation baselines, RL-ADR reduces power consumption by at least 14.5% and performs comparably to the oracle solution. Laixi Shi, Martin Hyungwoo Lee, Jaroslaw J. Sydir, Zhu Zhou, Yuejie Chi, Bin Li 0018 |
GLOBECOM | 6 |
| 2024 | Communication-Efficient Federated Optimization over Semi-Decentralized NetworksabstractIn large-scale federated and distributed learning, communication efficiency is one of the most challenging bottlenecks. While gossip communication—where agents can exchange information with their connected neighbors—is more cost-effective than communicating with the remote server, it often requires a greater number of communication rounds, especially for large and sparse networks. To tackle the trade-off, we examine the communication efficiency under a semi-decentralized communication protocol, in which agents can perform both agent-to-agent and agent-to-server communication in a probabilistic manner. We design a tailored communication-efficient algorithm over semi-decentralized networks, referred to as PISCO, which inherits the robustness to data heterogeneity thanks to gradient tracking and allows multiple local updates for saving communication. We establish the convergence rate of PISCO for nonconvex problems and show that PISCO enjoys a linear speedup in terms of the number of agents and local updates. Our numerical results highlight the superior communication efficiency of PISCO and its resilience to data heterogeneity and various network topologies. He Wang 0017, Yuejie Chi |
ICASSP | 2 |
| 2024 | Towards Non-Asymptotic Convergence for Diffusion-Based Generative ModelsabstractDiffusion models, which convert noise into new data instances by learning to reverse a Markov diffusion process, have become a cornerstone in contemporary generative modeling. While their practical power has now been widely recognized, the theoretical underpinnings remain far from mature. In this work, we develop a suite of non-asymptotic theory towards understanding the data generation process of diffusion models in discrete time, assuming access to $\ell_2$-accurate estimates of the (Stein) score functions. For a popular deterministic sampler (based on the probability flow ODE), we establish a convergence rate proportional to $1/T$ (with $T$ the total number of steps), improving upon past results; for another mainstream stochastic sampler (i.e., a type of the denoising diffusion probabilistic model), we derive a convergence rate proportional to $1/\sqrt{T}$, matching the state-of-the-art theory. Imposing only minimal assumptions on the target data distribution (e.g., no smoothness assumption is imposed), our results characterize how $\ell_2$ score estimation errors affect the quality of the data generation process. In contrast to prior works, our theory is developed based on an elementary yet versatile non-asymptotic approach without resorting to toolboxes for SDEs and ODEs. Gen Li 0005, Yuting Wei 0001, Yuxin Chen 0002, Yuejie Chi |
ICLR | 4 |
| 2024 | Accelerating Convergence of Score-Based Diffusion Models, ProvablyabstractScore-based diffusion models, while achieving remarkable empirical performance, often suffer from low sampling speed, due to extensive function evaluations needed during the sampling phase. Despite a flurry of recent activities towards speeding up diffusion generative modeling in practice, theoretical underpinnings for acceleration techniques remain severely limited. In this paper, we design novel training-free algorithms to accelerate popular deterministic (i.e., DDIM) and stochastic (i.e., DDPM) samplers. Our accelerated deterministic sampler converges at a rate $O(\frac{1}{{T}^2})$ with $T$ the number of steps, improving upon the $O(\frac{1}{T})$ rate for the DDIM sampler; and our accelerated stochastic sampler converges at a rate $O(\frac{1}{T})$, outperforming the rate $O(\frac{1}{\sqrt{T}})$ for the DDPM sampler. The design of our algorithms leverages insights from higher-order approximation, and shares similar intuitions as popular high-order ODE solvers like the DPM-Solver-2. Our theory accommodates $\ell_2$-accurate score estimates, and does not require log-concavity or smoothness on the target distribution. Gen Li 0005, Yu Huang 0023, Timofey Efimov, Yuting Wei 0001, Yuejie Chi, Yuxin Chen 0002 |
ICML | 5 |
| 2024 | Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM InferenceabstractMany computational factors limit broader deployment of large language models. In this paper, we focus on a memory bottleneck imposed by the key-value (KV) cache, a computational shortcut that requires storing previous KV pairs during decoding. While existing KV cache methods approach this problem by pruning or evicting large swaths of relatively less important KV pairs to dramatically reduce the memory footprint of the cache, they can have limited success in tasks that require recollecting a majority of previous tokens. To alleviate this issue, we propose LESS, a simple integration of a (nearly free) constant sized cache with eviction-based cache methods, such that all tokens can be queried at later decoding steps. Its ability to retain information throughout time shows merit on a variety of tasks where we demonstrate LESS can help reduce the performance gap from caching everything, sometimes even matching it, all while being efficient. Relevant code can be found at https://github.com/hdong920/LESS. Harry Dong, Xinyu Yang 0002, Zhenyu Zhang 0015, Zhangyang Wang, Yuejie Chi, Beidi Chen |
ICML | 5 |
| 2024 | Sample-Efficient Robust Multi-Agent Reinforcement Learning in the Face of Environmental UncertaintyabstractTo overcome the sim-to-real gap in reinforcement learning (RL), learned policies must maintain robustness against environmental uncertainties. While robust RL has been widely studied in single-agent regimes, in multi-agent environments, the problem remains understudied—despite the fact that the problems posed by environmental uncertainties are often exacerbated by strategic interactions. This work focuses on learning in distributionally robust Markov games (RMGs), a robust variant of standard Markov games, wherein each agent aims to learn a policy that maximizes its own worst-case performance when the deployed environment deviates within its own prescribed uncertainty set. This results in a set of robust equilibrium strategies for all agents that align with classic notions of game-theoretic equilibria. Assuming a non-adaptive sampling mechanism from a generative model, we propose a sample-efficient model-based algorithm (DRNVI) with finite-sample complexity guarantees for learning robust variants of various notions of game-theoretic equilibria. We also establish an information-theoretic lower bound for solving RMGs, which confirms the near-optimal sample complexity of DRNVI with respect to problem-dependent factors such as the size of the state space, the target accuracy, and the horizon length. Laixi Shi, Eric Mazumdar, Yuejie Chi, Adam Wierman |
ICML | 3 |
| 2024 | Federated Offline Reinforcement Learning: Collaborative Single-Policy Coverage SufficesabstractOffline reinforcement learning (RL), which seeks to learn an optimal policy using offline data, has garnered significant interest due to its potential in critical applications where online data collection is infeasible or expensive. This work explores the benefit of federated learning for offline RL, aiming at collaboratively leveraging offline datasets at multiple agents. Focusing on finite-horizon episodic tabular Markov decision processes (MDPs), we design FedLCB-Q, a variant of the popular model-free Q-learning algorithm tailored for federated offline RL. FedLCB-Q updates local Q-functions at agents with novel learning rate schedules and aggregates them at a central server using importance averaging and a carefully designed pessimistic penalty term. Our sample complexity analysis reveals that, with appropriately chosen parameters and synchronization schedules, FedLCB-Q achieves linear speedup in terms of the number of agents without requiring high-quality datasets at individual agents, as long as the local datasets collectively cover the state-action space visited by the optimal policy, highlighting the power of collaboration in the federated setting. In fact, the sample complexity almost matches that of the single-agent counterpart, as if all the data are stored at a central location, up to polynomial factors of the horizon length. Furthermore, FedLCB-Q is communication-efficient, where the number of communication rounds is only linear with respect to the horizon length up to logarithmic factors. Jiin Woo, Laixi Shi, Gauri Joshi, Yuejie Chi |
ICML | 4 |
| 2024 | Provably Robust Score-Based Diffusion Posterior Sampling for Plug-and-Play Image ReconstructionabstractIn a great number of tasks in science and engineering, the goal is to infer an unknown image from a small number of noisy measurements collected from a known forward model describing certain sensing or imaging modality. Due to resource constraints, this image reconstruction task is often extremely ill-posed, which necessitates the adoption of expressive prior information to regularize the solution space. Score-based diffusion models, thanks to its impressive empirical success, have emerged as an appealing candidate of an expressive prior in image reconstruction. In order to accommodate diverse tasks at once, it is of great interest to develop efficient, consistent and robust algorithms that incorporate unconditional score functions of an image prior distribution in conjunction with flexible choices of forward models.
This work develops an algorithmic framework for employing score-based diffusion models as an expressive data prior in nonlinear inverse problems with general forward models. Motivated by the plug-and-play framework in the imaging community, we introduce a diffusion plug-and-play method (DPnP) that alternatively calls two samplers, a proximal consistency sampler based solely on the likelihood function of the forward model, and a denoising diffusion sampler based solely on the score functions of the image prior. The key insight is that denoising under white Gaussian noise can be solved rigorously via both stochastic (i.e., DDPM-type) and deterministic (i.e., DDIM-type) samplers using the same set of score functions trained for generation. We establish both asymptotic and non-asymptotic performance guarantees of DPnP, and provide numerical experiments to illustrate its promise in solving both linear and nonlinear image reconstruction tasks. To the best of our knowledge, DPnP is the first provably-robust posterior sampling method for nonlinear inverse problems using unconditional diffusion priors. Xingyu Xu 0001, Yuejie Chi |
NeurIPS | 2 |
| 2024 | Learning Discrete Concepts in Latent Hierarchical ModelsabstractLearning concepts from natural high-dimensional data (e.g., images) holds potential in building human-aligned and interpretable machine learning models.
Despite its encouraging prospect, formalization and theoretical insights into this crucial task are still lacking.
In this work, we formalize concepts as discrete latent causal variables that are related via a hierarchical causal model that encodes different abstraction levels of concepts embedded in high-dimensional data (e.g., a dog breed and its eye shapes in natural images).
We formulate conditions to facilitate the identification of the proposed causal model, which reveals when learning such concepts from unsupervised data is possible.
Our conditions permit complex causal hierarchical structures beyond latent trees and multi-level directed acyclic graphs in prior work and can handle high-dimensional, continuous observed variables, which is well-suited for unstructured data modalities such as images.
We substantiate our theoretical claims with synthetic data experiments.
Further, we discuss our theory's implications for understanding the underlying mechanisms of latent diffusion models and provide corresponding empirical evidence for our theoretical insights. Guangyi Chen 0002, Biwei Huang, Eric P. Xing, Yuejie Chi, Kun Zhang 0001 |
NeurIPS | 5 |
| 2024 | The Sample-Communication Complexity Trade-off in Federated Q-LearningabstractWe consider the problem of Federated Q-learning, where $M$ agents aim to collaboratively learn the optimal Q-function of an unknown infinite horizon Markov Decision Process with finite state and action spaces. We investigate the trade-off between sample and communication complexity for the widely used class of intermittent communication algorithms. We first establish the converse result, where we show that any Federated Q-learning that offers a linear speedup with respect to number of agents in sample complexity needs to incur a communication cost of at least $\Omega(\frac{1}{1-\gamma})$, where $\gamma$ is the discount factor. We also propose a new Federated Q-learning algorithm, called Fed-DVR-Q, which is the first Federated Q-learning algorithm to simultaneously achieve order-optimal sample and communication complexities. Thus, together these results provide a complete characterization of the sample-communication complexity trade-off in Federated Q-learning. Sudeep Salgia, Yuejie Chi |
NeurIPS | 2 |
| 2024 | In-Context Learning with Representations: Contextual Generalization of Trained TransformersabstractIn-context learning (ICL) refers to a remarkable capability of pretrained large language models, which can learn a new task given a few examples during inference. However, theoretical understanding of ICL is largely under-explored, particularly whether transformers can be trained to generalize to unseen examples in a prompt, which will require the model to acquire contextual knowledge of the prompt for generalization. This paper investigates the training dynamics of transformers by gradient descent through the lens of non-linear regression tasks. The contextual generalization here can be attained via learning the template function for each task in-context, where all template functions lie in a linear space with $m$ basis functions. We analyze the training dynamics of one-layer multi-head transformers to {in-contextly} predict unlabeled inputs given partially labeled prompts, where the labels contain Gaussian noise and the number of examples in each prompt are not sufficient to determine the template. Under mild assumptions, we show that the training loss for a one-layer multi-head transformer converges linearly to a global minimum. Moreover, the transformer effectively learns to perform ridge regression over the basis functions. To our knowledge, this study is the first provable demonstration that transformers can learn contextual (i.e., template) information to generalize to both unseen examples and tasks when prompts contain only a small number of query-answer pairs. Tong Yang 0007, Yu Huang 0023, Yingbin Liang, Yuejie Chi |
NeurIPS | 4 |
| 2024 | Federated Natural Policy Gradient and Actor Critic Methods for Multi-task Reinforcement LearningabstractFederated reinforcement learning (RL) enables collaborative decision making of multiple distributed agents without sharing local data trajectories. In this work, we consider a multi-task setting, in which each agent has its own private reward function corresponding to different tasks, while sharing the same transition kernel of the environment. Focusing on infinite-horizon Markov decision processes, the goal is to learn a globally optimal policy that maximizes the sum of the discounted total rewards of all the agents in a decentralized manner, where each agent only communicates with its neighbors over some prescribed graph topology.
We develop federated vanilla and entropy-regularized natural policy gradient (NPG) methods in the tabular setting under softmax parameterization, where gradient tracking is applied to estimate the global Q-function to mitigate the impact of imperfect information sharing. We establish non-asymptotic global convergence guarantees under exact policy evaluation, where the rates are nearly independent of the size of the state-action space and illuminate the impacts of network size and connectivity. To the best of our knowledge, this is the first time that global convergence is established for federated multi-task RL using policy optimization. We further go beyond the tabular setting by proposing a federated natural actor critic (NAC) method for multi-task RL with function approximation, and establish its finite-time sample complexity taking the errors of function approximation into account. Tong Yang 0007, Shicong Cen, Yuting Wei 0001, Yuxin Chen 0002, Yuejie Chi |
NeurIPS | 5 |
| 2024 | Fast Policy Extragradient Methods for Competitive Games with Entropy RegularizationabstractThis paper investigates the problem of computing the equilibrium of competitive games in the form of two-player zero-sum games, which is often modeled as a constrained saddle-point optimization problem with probability simplex constraints. Despite recent efforts in understanding the last-iterate convergence of extragradient methods in the unconstrained setting, the theoretical underpinnings of these methods in the constrained settings, especially those using multiplicative updates, remain highly inadequate, even when the objective function is bilinear. Motivated by the algorithmic role of entropy regularization in single-agent reinforcement learning and game theory, we develop provably efficient extragradient methods to find the quantal response equilibrium (QRE)---which are solutions to zero-sum two-player matrix games with entropy regularization---at a linear rate. The proposed algorithms can be implemented in a decentralized manner, where each player executes symmetric and multiplicative updates iteratively using its own payoff without observing the opponent's actions directly. In addition, by controlling the knob of entropy regularization, the proposed algorithms can locate an approximate Nash equilibrium of the unregularized matrix game at a sublinear rate without assuming the Nash equilibrium to be unique. Our methods also lead to efficient policy extragradient algorithms for solving (entropy-regularized) zero-sum Markov games at similar rates. All of our convergence rates are nearly dimension-free, which are independent of the size of the state and action spaces up to logarithm factors, highlighting the positive role of entropy regularization for accelerating convergence. Shicong Cen, Yuting Wei 0001, Yuejie Chi |
J. Mach. Learn. Res. | 3 |
| 2024 | Distributionally Robust Model-Based Offline Reinforcement Learning with Near-Optimal Sample ComplexityabstractThis paper concerns the central issues of model robustness and sample efficiency in offline reinforcement learning (RL), which aims to learn to perform decision making from history data without active exploration. Due to uncertainties and variabilities of the environment, it is critical to learn a robust policy---with as few samples as possible---that performs well even when the deployed environment deviates from the nominal one used to collect the history dataset. We consider a distributionally robust formulation of offline RL, focusing on tabular robust Markov decision processes with an uncertainty set specified by the Kullback-Leibler divergence in both finite-horizon and infinite-horizon settings. To combat with sample scarcity, a model-based algorithm that combines distributionally robust value iteration with the principle of pessimism in the face of uncertainty is proposed, by penalizing the robust value estimates with a carefully designed data-driven penalty term. Under a mild and tailored assumption of the history dataset that measures distribution shift without requiring full coverage of the state-action space, we establish the finite-sample complexity of the proposed algorithms. We further develop an information-theoretic lower bound, which suggests that learning RMDPs is at least as hard as the standard MDPs when the uncertainty level is sufficient small, and corroborates the tightness of our upper bound up to polynomial factors of the (effective) horizon length for a range of uncertainty levels. To the best our knowledge, this provides the first provably near-optimal robust offline RL algorithm that learns under model uncertainty and partial coverage. Laixi Shi, Yuejie Chi |
J. Mach. Learn. Res. | 2 |
| 2024 | High-Probability Sample Complexities for Policy Evaluation With Linear Function ApproximationabstractThis paper is concerned with the problem of policy evaluation with linear function approximation in discounted infinite horizon Markov decision processes. We investigate the sample complexities required to guarantee a predefined estimation error of the best linear coefficients for two widely-used policy evaluation algorithms: the temporal difference (TD) learning algorithm and the two-timescale linear TD with gradient correction (TDC) algorithm. In both the on-policy setting, where observations are generated from the target policy, and the off-policy setting, where samples are drawn from a behavior policy potentially different from the target policy, we establish the first sample complexity bound with high-probability convergence guarantee that attains the optimal dependence on the tolerance level. We also exhibit an explicit dependence on problem-related quantities, and show in the on-policy setting that our upper bound matches the minimax lower bound on crucial problem parameters, including the choice of the feature map and the problem dimension. Gen Li 0005, Weichen Wu, Yuejie Chi, Cong Ma 0001, Alessandro Rinaldo, Yuting Wei 0001 |
IEEE Trans. Inf. Theory | 3 |
| 2023 | Understanding Masked Autoencoders via Hierarchical Latent Variable ModelsabstractMasked autoencoder (MAE), a simple and effective self-supervised learning framework based on the reconstruction of masked image regions, has recently achieved prominent success in a variety of vision tasks. Despite the emergence of intriguing empirical observations on MAE, a theoretically principled understanding is still lacking. In this work, we formally characterize and justify existing empirical in-sights and provide theoretical guarantees of MAE. We formulate the underlying data-generating process as a hierarchical latent variable model, and show that under reasonable assumptions, MAE provably identifies a set of latent variables in the hierarchical model, explaining why MAE can extract high-level information from pixels. Further, we show how key hyperparameters in MAE (the masking ratio and the patch size) determine which true latent variables to be recovered, therefore influencing the level of semantic information in the representation. Specifically, extremely large or small masking ratios inevitably lead to low-level representations. Our theory offers coherent explanations of existing empirical observations and provides insights for potential empirical improvements and fundamental limitations of the masked-reconstruction paradigm. We conduct extensive experiments to validate our theoretical insights. Martin Q. Ma, Guangyi Chen 0002, Eric P. Xing, Yuejie Chi, Louis-Philippe Morency, Kun Zhang 0001 |
CVPR | 5 |
| 2023 | Deep Unfolded Tensor Robust PCA With Self-Supervised LearningabstractTensor robust principal component analysis (RPCA), which seeks to separate a low-rank tensor from its sparse corruptions, has been crucial in data science and machine learning where tensor structures are becoming more prevalent. While powerful, existing tensor RPCA algorithms can be difficult to use in practice, as their performance can be sensitive to the choice of additional hyperparameters, which are not straightforward to tune. In this paper, we describe a fast and simple self-supervised model for tensor RPCA using deep unfolding by only learning four hyperparameters. Despite its simplicity, our model expunges the need for ground truth labels while maintaining competitive or even greater performance compared to supervised deep unfolding. Furthermore, our model is capable of operating in extreme data-starved scenarios. We demonstrate these claims on a mix of synthetic data and real-world tasks, comparing performance against previously studied supervised deep unfolding methods and Bayesian optimization baselines. Harry Dong, Megna Shah, Sean Donegan, Yuejie Chi |
ICASSP | 4 |
| 2023 | Asynchronous Gradient Play in Zero-Sum Multi-agent Games
Ruicheng Ao, Shicong Cen, Yuejie Chi |
ICLR | 3 |
| 2023 | Faster Last-iterate Convergence of Policy Optimization in Zero-Sum Markov Games
Shicong Cen, Yuejie Chi, Simon S. Du |
ICLR | 2 |
| 2023 | The Power of Preconditioning in Overparameterized Low-Rank Matrix SensingabstractWe propose $\textsf{ScaledGD($\lambda$)}$, a preconditioned gradient descent method to tackle the low-rank matrix sensing problem when the true rank is unknown, and when the matrix is possibly ill-conditioned. Using overparametrized factor representations, $\textsf{ScaledGD($\lambda$)}$ starts from a small random initialization, and proceeds by gradient descent with a specific form of preconditioning with a fixed damping term to combat overparameterization. At the expense of light computational overhead incurred by preconditioners, $\textsf{ScaledGD($\lambda$)}$ is remarkably robust to ill-conditioning compared to vanilla gradient descent ($\mathsf{GD}$). Specifically, we show that, under the Gaussian design, $\textsf{ScaledGD($\lambda$)}$ converges to the true low-rank matrix at a constant linear rate that is independent of the condition number (apart from a short nearly dimension-free burdening period), with near-optimal sample complexity. This significantly improves upon the convergence rate of vanilla $\mathsf{GD}$ which suffers from a polynomial dependency with the condition number. Our work provides evidence on the power of preconditioning in accelerating the convergence without hurting generalization in overparameterized learning. Xingyu Xu 0001, Yandi Shen, Yuejie Chi, Cong Ma 0001 |
ICML | 3 |
| 2023 | The Blessing of Heterogeneity in Federated Q-Learning: Linear Speedup and BeyondabstractIn this paper, we consider federated Q-learning, which aims to learn an optimal Q-function by periodically aggregating local Q-estimates trained on local data alone. Focusing on infinite-horizon tabular Markov decision processes, we provide sample complexity guarantees for both the synchronous and asynchronous variants of federated Q-learning. In both cases, our bounds exhibit a linear speedup with respect to the number of agents and sharper dependencies on other salient problem parameters. Moreover, existing approaches to federated Q-learning adopt an equally-weighted averaging of local Q-estimates, which can be highly sub-optimal in the asynchronous setting since the local trajectories can be highly heterogeneous due to different local behavior policies. Existing sample complexity scales inverse proportionally to the minimum entry of the stationary state-action occupancy distributions over all agents, requiring that every agent covers the entire state-action space. Instead, we propose a novel importance averaging algorithm, giving larger weights to more frequently visited state-action pairs. The improved sample complexity scales inverse proportionally to the minimum entry of the average stationary state-action occupancy distribution of all agents, thus only requiring the agents collectively cover the entire state-action space, unveiling the blessing of heterogeneity. Jiin Woo, Gauri Joshi, Yuejie Chi |
ICML | 3 |
| 2023 | Reward-agnostic Fine-tuning: Provable Statistical Benefits of Hybrid Reinforcement LearningabstractThis paper studies tabular reinforcement learning (RL) in the hybrid setting, which assumes access to both an offline dataset and online interactions with the unknown environment. A central question boils down to how to efficiently utilize online data to strengthen and complement the offline dataset and enable effective policy fine-tuning. Leveraging recent advances in reward-agnostic exploration and offline RL, we design a three-stage hybrid RL algorithm that beats the best of both worlds --- pure offline RL and pure online RL --- in terms of sample complexities. The proposed algorithm does not require any reward information during data collection. Our theory is developed based on a new notion called **single-policy partial concentrability**, which captures the trade-off between distribution mismatch and miscoverage and guides the interplay between offline and online data. Gen Li 0005, Wenhao Zhan, Jason D. Lee, Yuejie Chi, Yuxin Chen 0002 |
NeurIPS | 4 |
| 2023 | Seeing is not Believing: Robust Reinforcement Learning against Spurious CorrelationabstractRobustness has been extensively studied in reinforcement learning (RL) to handle various forms of uncertainty such as random perturbations, rare events, and malicious attacks. In this work, we consider one critical type of robustness against spurious correlation, where different portions of the state do not have correlations induced by unobserved confounders. These spurious correlations are ubiquitous in real-world tasks, for instance, a self-driving car usually observes heavy traffic in the daytime and light traffic at night due to unobservable human activity. A model that learns such useless or even harmful correlation could catastrophically fail when the confounder in the test case deviates from the training one. Although motivated, enabling robustness against spurious correlation poses significant challenges since the uncertainty set, shaped by the unobserved confounder and causal structure, is difficult to characterize and identify. Existing robust algorithms that assume simple and unstructured uncertainty sets are therefore inadequate to address this challenge. To solve this issue, we propose Robust State-Confounded Markov Decision Processes (RSC-MDPs) and theoretically demonstrate its superiority in avoiding learning spurious correlations compared with other robust RL counterparts. We also design an empirical algorithm to learn the robust optimal policy for RSC-MDPs, which outperforms all baselines in eight realistic self-driving and manipulation tasks. Wenhao Ding, Laixi Shi, Yuejie Chi, Ding Zhao |
NeurIPS | 3 |
| 2023 | Identification of Nonlinear Latent Hierarchical ModelsabstractIdentifying latent variables and causal structures from observational data is essential to many real-world applications involving biological data, medical data, and unstructured data such as images and languages. However, this task can be highly challenging, especially when observed variables are generated by causally related latent variables and the relationships are nonlinear.
In this work, we investigate the identification problem for nonlinear latent hierarchical causal models in which observed variables are generated by a set of causally related latent variables, and some latent variables may not have observed children. We show that the identifiability of causal structures and latent variables (up to invertible transformations) can be achieved under mild assumptions: on causal structures, we allow for multiple paths between any pair of variables in the graph, which relaxes latent tree assumptions in prior work; on structural functions, we permit general nonlinearity and multi-dimensional continuous variables, alleviating existing work's parametric assumptions. Specifically, we first develop an identification criterion in the form of novel identifiability guarantees for an elementary latent variable model. Leveraging this criterion, we show that both causal structures and latent variables of the hierarchical model can be identified asymptotically by explicitly constructing an estimation procedure. To the best of our knowledge, our work is the first to establish identifiability guarantees for both causal structures and latent variables in nonlinear latent hierarchical models. Biwei Huang, Feng Xie 0002, Eric P. Xing, Yuejie Chi, Kun Zhang 0001 |
NeurIPS | 5 |
| 2023 | The Curious Price of Distributional Robustness in Reinforcement Learning with a Generative ModelabstractThis paper investigates model robustness in reinforcement learning (RL) via the framework of distributionally robust Markov decision processes (RMDPs). Despite recent efforts, the sample complexity of RMDPs is much less understood regardless of the uncertainty set in use; in particular, there exist large gaps between existing upper and lower bounds, and it is unclear if distributional robustness bears any statistical implications when benchmarked against standard RL. In this paper, assuming access to a generative model, we derive the sample complexity of RMDPs---when the uncertainty set is measured via either total variation or $\chi^2$ divergence over the full range of uncertainty levels---using a model-based algorithm called distributionally robust value iteration, and develop minimax lower bounds to benchmark its tightness. Our results not only strengthen the prior art in both directions of upper and lower bounds, but also deliver surprising messages that learning RMDPs is not necessarily easier or more difficult than standard MDPs. In the case of total variation, we establish the minimax-optimal sample complexity of RMDPs which is always smaller than that of standard MDPs. In the case of $\chi^2$ divergence, we establish the sample complexity of RMDPs that is tight up to polynomial factors of the effective horizon, and grows linearly with respect to the uncertainty level when it approaches infinity. Laixi Shi, Gen Li 0005, Yuting Wei 0001, Yuxin Chen 0002, Matthieu Geist, Yuejie Chi |
NeurIPS | 6 |
| 2023 | Counterfactual Generation with Identifiability GuaranteesabstractCounterfactual generation lies at the core of various machine learning tasks, including image translation and controllable text generation. This generation process usually requires the identification of the disentangled latent representations, such as content and style, that underlie the observed data. However, it becomes more challenging when faced with a scarcity of paired data and labelling information. Existing disentangled methods crucially rely on oversimplified assumptions, such as assuming independent content and style variables, to identify the latent variables, even though such assumptions may not hold for complex data distributions. For instance, food reviews tend to involve words like “tasty”, whereas movie reviews commonly contain words such as “thrilling” for the same positive sentiment. This problem is exacerbated when data are sampled from multiple domains since the dependence between content and style may vary significantly over domains. In this work, we tackle the domain-varying dependence between the content and the style variables inherent in the counterfactual generation task. We provide identification guarantees for such latent-variable models by leveraging the relative sparsity of the influences from different latent variables. Our theoretical insights enable the development of a doMain AdapTive counTerfactual gEneration model, called (MATTE). Our theoretically grounded framework achieves state-of-the-art performance in unsupervised style transfer tasks, where neither paired data nor style labels are utilized, across four large-scale datasets. Hanqi Yan, Lin Gui 0003, Yuejie Chi, Eric P. Xing, Yulan He 0001, Kun Zhang 0001 |
NeurIPS | 4 |
| 2023 | Offline Reinforcement Learning with On-Policy Q-Function Regularization
Laixi Shi, Robert Dadashi, Yuejie Chi, Pablo Samuel Castro, Matthieu Geist |
ECML/PKDD (4) | 3 |
| 2023 | A trajectory is worth three sentences: multimodal transformer for offline reinforcement learningabstractTransformers hold tremendous promise in solving offline reinforcement learning (RL) by formulating it as a sequence modeling problem inspired by language modeling (LM). Prior works using transformers model a sample (trajectory) of RL as one sequence analogous to a sequence of words (one sentence) in LM, despite the fact that each trajectory includes tokens from three diverse modalities: state, action, and reward, while a sentence contains words only. Rather than taking a modality-agnostic approach which uniformly models the tokens from different modalities as one sequence, we propose a multimodal sequence modeling approach in which a trajectory (one “sentence”) of three modalities (state, action, reward) is disentangled into three unimodal ones (three “sentences”). We investigate the correlation of different modalities during sequential decision-making and use the insights to design a multimodal transformer, named Decision Transducer (DTd). DTd outperforms prior art in offline RL on the conducted D4RL benchmarks and enjoys better sample efficiency and algorithm flexibility. Our code is made publicly here. Mengdi Xu, Laixi Shi, Yuejie Chi |
UAI | 4 |
| 2022 | Batch Active Learning with Graph Neural Networks via Multi-Agent Deep Reinforcement LearningabstractGraph neural networks (GNNs) have achieved tremendous success in many graph learning tasks such as node classification, graph classification and link prediction. For the classification task, GNNs' performance often highly depends on the number of labeled nodes and thus could be significantly hampered due to the expensive annotation cost. The sparse literature on active learning for GNNs has primarily focused on selecting only one sample each iteration, which becomes inefficient for large scale datasets. In this paper, we study the batch active learning setting for GNNs where the learning agent can acquire labels of multiple samples at each time. We formulate batch active learning as a cooperative multi-agent reinforcement learning problem and present a novel reinforced batch-mode active learning framework BiGeNe. To avoid the combinatorial explosion of the joint action space, we introduce a value decomposition method that factorizes the total Q-value into the average of individual Q-values. Moreover, we propose a novel multi-agent Q-network consisting of a graph convolutional network (GCN) component and a gated recurrent unit (GRU) component. The GCN component takes both the informativeness and inter-dependences between nodes into account and the GRU component enables the agent to consider interactions between selected nodes in the same batch. Experimental results on multiple public datasets demonstrate the effectiveness and efficiency of our proposed method. Hanghang Tong, Yinglong Xia, Yuejie Chi, Lei Ying 0001 |
AAAI | 5 |
| 2022 | Scaling and Scalability: Provable Nonconvex Low-Rank Tensor CompletionabstractTensors, which provide a powerful and flexible model for representing multi-attribute data and multi-way interactions, play an indispensable role in modern data science across various fields in science and engineering. A fundamental task is tensor completion, which aims to faithfully recover the tensor from a small subset of its entries in a statistically and computationally efficient manner. Harnessing the low-rank structure of tensors in the Tucker decomposition, this paper develops a scaled gradient descent (ScaledGD) algorithm to directly recover the tensor factors with tailored spectral initializations, and shows that it provably converges at a linear rate independent of the condition number of the ground truth tensor for tensor completion as soon as the sample size is above the order of $n^{3/2}$ ignoring other parameter dependencies, where $n$ is the dimension of the tensor. To the best of our knowledge, ScaledGD is the first algorithm that achieves near-optimal statistical and computational complexities simultaneously for low-rank tensor completion with the Tucker decomposition. Our algorithm highlights the power of appropriate preconditioning in accelerating nonconvex statistical estimation, where the iteration-varying preconditioners promote desirable invariance properties of the trajectory with respect to the underlying symmetry in low-rank tensor factorization. Tian Tong, Cong Ma 0001, Ashley Prater-Bennette, Erin E. Tripp, Yuejie Chi |
AISTATS | 5 |
| 2022 | Privacy-Preserving Federated Multi-Task Linear Regression: A One-Shot Linear Mixing Approach Inspired By Graph RegularizationabstractWe investigate multi-task learning (MTL), where multiple learning tasks are performed jointly rather than separately to leverage their similarities and improve performance. We focus on the federated multi-task linear regression setting, where each machine possesses its own data for individual tasks and sharing the full local data between machines is prohibited. Motivated by graph regularization, we propose a novel fusion framework that only requires a one-shot communication of local estimates. Our method linearly combines the local estimates to produce an improved estimate for each task, and we show that the ideal mixing weight for fusion is a function of task similarity and task difficulty. A practical algorithm is developed and shown to significantly reduce mean squared error (MSE) on synthetic data, as well as improve performance on an income prediction task where the real-world data is disaggregated by race. Harlin Lee, Andrea L. Bertozzi, Jelena Kovacevic, Yuejie Chi |
ICASSP | 4 |
| 2022 | Accelerating ILL-Conditioned Robust Low-Rank Tensor RegressionabstractAn important problem that arises across different applications in signal processing, machine learning, and data science is to reliably estimate a tensor from a small number of measurements that are possibly corrupted. Leveraging the low-rank structure under the Tucker decomposition, we propose a provably efficient algorithm that directly estimates the tensor factors by solving a nonsmooth and nonconvex composite optimization problem that minimizes the least absolute deviation loss. The proposed algorithm—built on subgradient methods—harnesses preconditioners that are designed to be equivariant w.r.t. the low-rank parameterization, and is shown to achieve local linear convergence at a constant rate under the Gaussian design. Numerical experiments are provided to corroborate the superior performance of the proposed algorithm. Tian Tong, Cong Ma 0001, Yuejie Chi |
ICASSP | 3 |
| 2022 | Active Heterogeneous Graph Neural Networks with Per-step Meta-Q-LearningabstractRecent years have witnessed the superior performance of heterogeneous graph neural networks (HGNNs) in dealing with heterogeneous information networks (HINs). Nonetheless, the success of HGNNs often depends on the availability of sufficient labeled training data, which can be very expensive to obtain in real scenarios. Active learning provides an effective solution to tackle the data scarcity challenge. For the vast majority of the existing work regarding active learning on graphs, they mainly focus on homogeneous graphs, and thus fall in short or even become inapplicable on HINs. In this paper, we study the active learning problem with HGNNs and propose a novel meta-reinforced active learning framework MetRA. Previous reinforced active learning algorithms train the policy network on labeled source graphs and directly transfer the policy to the target graph without any adaptation. To better exploit the information from the target graph in the adaptation phase, we propose a novel policy transfer algorithm based on meta-Q-learning termed per-step MQL. Empirical evaluations on HINs demonstrate the effectiveness of our proposed framework. The improvement over the best baseline is up to 7% in Micro-F1. Yinglong Xia, Yuejie Chi, Lei Ying 0001, Hanghang Tong |
ICDM | 4 |
| 2022 | Pessimistic Q-Learning for Offline Reinforcement Learning: Towards Optimal Sample ComplexityabstractOffline or batch reinforcement learning seeks to learn a near-optimal policy using history data without active exploration of the environment. To counter the insufficient coverage and sample scarcity of many offline datasets, the principle of pessimism has been recently introduced to mitigate high bias of the estimated values. While pessimistic variants of model-based algorithms (e.g., value iteration with lower confidence bounds) have been theoretically investigated, their model-free counterparts — which do not require explicit model estimation — have not been adequately studied, especially in terms of sample efficiency. To address this inadequacy, we study a pessimistic variant of Q-learning in the context of finite-horizon Markov decision processes, and characterize its sample complexity under the single-policy concentrability assumption which does not require the full coverage of the state-action space. In addition, a variance-reduced pessimistic Q-learning algorithm is proposed to achieve near-optimal sample complexity. Altogether, this work highlights the efficiency of model-free algorithms in offline RL when used in conjunction with pessimism and variance reduction. Laixi Shi, Gen Li 0005, Yuting Wei 0001, Yuxin Chen 0002, Yuejie Chi |
ICML | 5 |
| 2022 | Minimax-Optimal Multi-Agent RL in Markov Games With a Generative ModelabstractThis paper studies multi-agent reinforcement learning in Markov games, with the goal of learning Nash equilibria or coarse correlated equilibria (CCE) sample-optimally. All prior results suffer from at least one of the two obstacles: the curse of multiple agents and the barrier of long horizon, regardless of the sampling protocol in use. We take a step towards settling this problem, assuming access to a flexible sampling mechanism: the generative model. Focusing on non-stationary finite-horizon Markov games, we develop a fast learning algorithm called Q-FTRL and an adaptive sampling scheme that leverage the optimism principle in online adversarial learning (particularly the Follow-the-Regularized-Leader (FTRL) method). Our algorithm learns an $\varepsilon$-approximate CCE in a general-sum Markov game using $$ \widetilde{O}\bigg( \frac{H^4 S \sum_{i=1}^m A_i}{\varepsilon^2} \bigg) $$ samples, where $m$ is the number of players, $S$ indicates the number of states, $H$ is the horizon, and $A_i$ denotes the number of actions for the $i$-th player. This is minimax-optimal (up to log factor) when $m$ is fixed. When applied to two-player zero-sum Markov games, our algorithm provably finds an $\varepsilon$-approximate Nash equilibrium with a minimal number of samples. Along the way, we derive a refined regret bound for FTRL that makes explicit the role of variance-type quantities, which might be of independent interest. Gen Li 0005, Yuejie Chi, Yuting Wei 0001, Yuxin Chen 0002 |
NeurIPS | 2 |
| 2022 | SoteriaFL: A Unified Framework for Private Federated Learning with Communication CompressionabstractTo enable large-scale machine learning in bandwidth-hungry environments such as wireless networks, significant progress has been made recently in designing communication-efficient federated learning algorithms with the aid of communication compression. On the other end, privacy preserving, especially at the client level, is another important desideratum that has not been addressed simultaneously in the presence of advanced communication compression techniques yet. In this paper, we propose a unified framework that enhances the communication efficiency of private federated learning with communication compression. Exploiting both general compression operators and local differential privacy, we first examine a simple algorithm that applies compression directly to differentially-private stochastic gradient descent, and identify its limitations. We then propose a unified framework SoteriaFL for private federated learning, which accommodates a general family of local gradient estimators including popular stochastic variance-reduced gradient methods and the state-of-the-art shifted compression scheme. We provide a comprehensive characterization of its performance trade-offs in terms of privacy, utility, and communication complexity, where SoteriaFL is shown to achieve better communication complexity without sacrificing privacy nor utility than other private federated learning algorithms without communication compression. Zhize Li 0001, Boyue Li, Yuejie Chi |
NeurIPS | 4 |
| 2022 | BEER: Fast $O(1/T)$ Rate for Decentralized Nonconvex Optimization with Communication CompressionabstractCommunication efficiency has been widely recognized as the bottleneck for large-scale decentralized machine learning applications in multi-agent or federated environments. To tackle the communication bottleneck, there have been many efforts to design communication-compressed algorithms for decentralized nonconvex optimization, where the clients are only allowed to communicate a small amount of quantized information (aka bits) with their neighbors over a predefined graph topology. Despite significant efforts, the state-of-the-art algorithm in the nonconvex setting still suffers from a slower rate of convergence $O((G/T)^{2/3})$ compared with their uncompressed counterpart, where $G$ measures the data heterogeneity across different clients, and $T$ is the number of communication rounds. This paper proposes BEER, which adopts communication compression with gradient tracking, and shows it converges at a faster rate of $O(1/T)$. This significantly improves over the state-of-the-art rate, by matching the rate without compression even under arbitrary data heterogeneity. Numerical experiments are also provided to corroborate our theory and confirm the practical superiority of beer in the data heterogeneous regime. Boyue Li, Zhize Li 0001, Peter Richtárik, Yuejie Chi |
NeurIPS | 5 |
| 2022 | Scaling and Scalability: Provable Nonconvex Low-Rank Tensor Estimation from Incomplete MeasurementsabstractTensors, which provide a powerful and flexible model for representing multi-attribute data and multi-way interactions, play an indispensable role in modern data science across various fields in science and engineering. A fundamental task is to faithfully recover the tensor from highly incomplete measurements in a statistically and computationally efficient manner. Harnessing the low-rank structure of tensors in the Tucker decomposition, this paper develops a scaled gradient descent (ScaledGD) algorithm to directly recover the tensor factors with tailored spectral initializations, and shows that it provably converges at a linear rate independent of the condition number of the ground truth tensor for two canonical problems --- tensor completion and tensor regression --- as soon as the sample size is above the order of $n^{3/2}$ ignoring other parameter dependencies, where $n$ is the dimension of the tensor. This leads to an extremely scalable approach to low-rank tensor estimation compared with prior art, which suffers from at least one of the following drawbacks: extreme sensitivity to ill-conditioning, high per-iteration costs in terms of memory and computation, or poor sample complexity guarantees. To the best of our knowledge, ScaledGD is the first algorithm that achieves near-optimal statistical and computational complexities simultaneously for low-rank tensor completion with the Tucker decomposition. Our algorithm highlights the power of appropriate preconditioning in accelerating nonconvex statistical estimation, where the iteration-varying preconditioners promote desirable invariance properties of the trajectory with respect to the underlying symmetry in low-rank tensor factorization. Tian Tong, Cong Ma 0001, Ashley Prater-Bennette, Erin E. Tripp, Yuejie Chi |
J. Mach. Learn. Res. | 5 |
| 2022 | Sample Complexity of Asynchronous Q-Learning: Sharper Analysis and Variance ReductionabstractAsynchronous Q-learning aims to learn the optimal action-value function (or Q-function) of a Markov decision process (MDP), based on a single trajectory of Markovian samples induced by a behavior policy. Focusing on a$\gamma $-discounted MDP with state space$\mathcal {S}$and action space$\mathcal {A}$, we demonstrate that the$\ell _{\infty }$-based sample complexity of classical asynchronous Q-learning — namely, the number of samples needed to yield an entrywise$\varepsilon $-accurate estimate of the Q-function — is at most on the order of$\frac {1}{ \mu _{\mathsf {min}}(1-\gamma)^{5}\varepsilon ^{2}}+ \frac { t_{\mathsf {mix}}}{ \mu _{\mathsf {min}}(1-\gamma)}$up to some logarithmic factor, provided that a proper constant learning rate is adopted. Here,$t_{\mathsf {mix}}$and$\mu _{\mathsf {min}}$denote respectively the mixing time and the minimum state-action occupancy probability of the sample trajectory. The first term of this bound matches the sample complexity in the synchronous case with independent samples drawn from the stationary distribution of the trajectory. The second term reflects the cost taken for the empirical distribution of the Markovian trajectory to reach a steady state, which is incurred at the very beginning and becomes amortized as the algorithm runs. Encouragingly, the above bound improves upon the state-of-the-art result by a factor of at least$|\mathcal {S}||\mathcal {A}|$for all scenarios, and by a factor of at least$t_{\mathsf {mix}}|\mathcal {S}||\mathcal {A}|$for any sufficiently small accuracy level$\varepsilon $. Further, we demonstrate that the scaling on the effective horizon$\frac {1}{1-\gamma }$can be improved by means of variance reduction. Gen Li 0005, Yuting Wei 0001, Yuejie Chi, Yuantao Gu, Yuxin Chen 0002 |
IEEE Trans. Inf. Theory | 3 |
| 2021 | Softmax Policy Gradient Methods Can Take Exponential Time to ConvergeabstractThe softmax policy gradient (PG) method, which performs gradient ascent under softmax policy parameterization, is arguably one of the de facto implementations of policy optimization in modern reinforcement learning. For $\gamma$-discounted infinite-horizon tabular Markov decision processes (MDPs), remarkable progress has recently been achieved towards establishing global convergence of softmax PG methods in finding a near-optimal policy. However, prior results fall short of delineating clear dependencies of convergence rates on salient parameters such as the cardinality of the state space $\mathcal{S}$ and the effective horizon $\frac{1}{1-\gamma}$, both of which could be excessively large. In this paper, we deliver a pessimistic message regarding the iteration complexity of softmax PG methods, despite assuming access to exact gradient computation. Specifically, we demonstrate that the softmax PG method with stepsize $\eta$ can take \[ \frac{1}{\eta} |\mathcal{S}|^{2^{\Omega\big(\frac{1}{1-\gamma}\big)}} \text{iterations} \]{to} converge, even in the presence of a benign policy initialization and an initial state distribution amenable to exploration (so that the distribution mismatch coefficient is not exceedingly large). This is accomplished by characterizing the algorithmic dynamics over a carefully-constructed MDP containing only three actions. Our exponential lower bound hints at the necessity of carefully adjusting update rules or enforcing proper regularization in accelerating PG methods. Gen Li 0005, Yuting Wei 0001, Yuejie Chi, Yuantao Gu, Yuxin Chen 0002 |
COLT | 3 |
| 2021 | Plug-And-Play Image Reconstruction Meets Stochastic Variance-Reduced Gradient MethodsabstractPlug-and-play (PnP) methods have recently emerged as a powerful framework for image reconstruction that can flexibly combine different physics-based observation models with data-driven image priors in the form of denoisers, and achieve state-of-the-art image reconstruction quality in many applications. In this paper, we aim to further improve the computational efficacy of PnP methods by designing a new algorithm that makes use of stochastic variance-reduced gradients (SVRG), a nascent idea to accelerate runtime in stochastic optimization. Compared with existing PnP methods using batch gradients or stochastic gradients, the new algorithm, called PnP-SVRG, achieves comparable or better accuracy of image reconstruction at a much faster computational speed. Extensive numerical experiments are provided to demonstrate the benefits of the proposed algorithm through the application of compressive imaging using partial Fourier measurements in conjunction with a wide variety of popular image denoisers. Vincent Monardo, Abhiram Iyer, Sean Donegan, Marc De Graef, Yuejie Chi |
ICIP | 5 |
| 2021 | Tightening the Dependence on Horizon in the Sample Complexity of Q-LearningabstractQ-learning, which seeks to learn the optimal Q-function of a Markov decision process (MDP) in a model-free fashion, lies at the heart of reinforcement learning. Focusing on the synchronous setting (such that independent samples for all state-action pairs are queried via a generative model in each iteration), substantial progress has been made recently towards understanding the sample efficiency of Q-learning. To yield an entrywise $\varepsilon$-accurate estimate of the optimal Q-function, state-of-the-art theory requires at least an order of $\frac{|S||A|}{(1-\gamma)^5\varepsilon^{2}}$ samples in the infinite-horizon $\gamma$-discounted setting. In this work, we sharpen the sample complexity of synchronous Q-learning to the order of $\frac{|S||A|}{(1-\gamma)^4\varepsilon^2}$ (up to some logarithmic factor) for any $0<\varepsilon <1$, leading to an order-wise improvement in $\frac{1}{1-\gamma}$. Analogous results are derived for finite-horizon MDPs as well. Notably, our sample complexity analysis unveils the effectiveness of vanilla Q-learning, which matches that of speedy Q-learning without requiring extra computation and storage. Our result is obtained by identifying novel error decompositions and recursion relations, which might shed light on how to study other variants of Q-learning. Gen Li 0005, Changxiao Cai, Yuxin Chen 0002, Yuantao Gu, Yuting Wei 0001, Yuejie Chi |
ICML | 6 |
| 2021 | Fast Policy Extragradient Methods for Competitive Games with Entropy RegularizationabstractThis paper investigates the problem of computing the equilibrium of competitive games, which is often modeled as a constrained saddle-point optimization problem with probability simplex constraints. Despite recent efforts in understanding the last-iterate convergence of extragradient methods in the unconstrained setting, the theoretical underpinnings of these methods in the constrained settings, especially those using multiplicative updates, remain highly inadequate, even when the objective function is bilinear. Motivated by the algorithmic role of entropy regularization in single-agent reinforcement learning and game theory, we develop provably efficient extragradient methods to find the quantal response equilibrium (QRE)---which are solutions to zero-sum two-player matrix games with entropy regularization---at a linear rate. The proposed algorithms can be implemented in a decentralized manner, where each player executes symmetric and multiplicative updates iteratively using its own payoff without observing the opponent's actions directly. In addition, by controlling the knob of entropy regularization, the proposed algorithms can locate an approximate Nash equilibrium of the unregularized matrix game at a sublinear rate without assuming the Nash equilibrium to be unique. Our methods also lead to efficient policy extragradient algorithms for solving entropy-regularized zero-sum Markov games at a linear rate. All of our convergence rates are nearly dimension-free, which are independent of the size of the state and action spaces up to logarithm factors, highlighting the positive role of entropy regularization for accelerating convergence. Shicong Cen, Yuting Wei 0001, Yuejie Chi |
NeurIPS | 3 |
| 2021 | Sample-Efficient Reinforcement Learning Is Feasible for Linearly Realizable MDPs with Limited RevisitingabstractLow-complexity models such as linear function representation play a pivotal role in enabling sample-efficient reinforcement learning (RL). The current paper pertains to a scenario with value-based linear representation, which postulates linear realizability of the optimal Q-function (also called the ``linear $Q^{\star}$ problem''). While linear realizability alone does not allow for sample-efficient solutions in general, the presence of a large sub-optimality gap is a potential game changer, depending on the sampling mechanism in use. Informally, sample efficiency is achievable with a large sub-optimality gap when a generative model is available, but is unfortunately infeasible when we turn to standard online RL settings. We make progress towards understanding this linear $Q^{\star}$ problem by investigating a new sampling protocol, which draws samples in an online/exploratory fashion but allows one to backtrack and revisit previous states. This protocol is more flexible than the standard online RL setting, while being practically relevant and far more restrictive than the generative model. We develop an algorithm tailored to this setting, achieving a sample complexity that scales polynomially with the feature dimension, the horizon, and the inverse sub-optimality gap, but not the size of the state/action space. Our findings underscore the fundamental interplay between sampling protocols and low-complexity function representation in RL. Gen Li 0005, Yuxin Chen 0002, Yuejie Chi, Yuantao Gu, Yuting Wei 0001 |
NeurIPS | 3 |
| 2021 | Breaking the Sample Complexity Barrier to Regret-Optimal Model-Free Reinforcement LearningabstractAchieving sample efficiency in online episodic reinforcement learning (RL) requires optimally balancing exploration and exploitation. When it comes to a finite-horizon episodic Markov decision process with $S$ states, $A$ actions and horizon length $H$, substantial progress has been achieved towards characterizing the minimax-optimal regret, which scales on the order of $\sqrt{H^2SAT}$ (modulo log factors) with $T$ the total number of samples. While several competing solution paradigms have been proposed to minimize regret, they are either memory-inefficient, or fall short of optimality unless the sample size exceeds an enormous threshold (e.g., $S^6A^4 \,\mathrm{poly}(H)$ for existing model-free methods).To overcome such a large sample size barrier to efficient RL, we design a novel model-free algorithm, with space complexity $O(SAH)$, that achieves near-optimal regret as soon as the sample size exceeds the order of $SA\,\mathrm{poly}(H)$. In terms of this sample size requirement (also referred to the initial burn-in cost), our method improves --- by at least a factor of $S^5A^3$ --- upon any prior memory-efficient algorithm that is asymptotically regret-optimal. Leveraging the recently introduced variance reduction strategy (also called {\em reference-advantage decomposition}), the proposed algorithm employs an {\em early-settled} reference update rule, with the aid of two Q-learning sequences with upper and lower confidence bounds. The design principle of our early-settled variance reduction method might be of independent interest to other RL settings that involve intricate exploration-exploitation trade-offs. Gen Li 0005, Laixi Shi, Yuxin Chen 0002, Yuantao Gu, Yuejie Chi |
NeurIPS | 5 |
| 2021 | Accelerating Ill-Conditioned Low-Rank Matrix Estimation via Scaled Gradient DescentabstractLow-rank matrix estimation is a canonical problem that finds numerous applications in signal processing, machine learning and imaging science. A popular approach in practice is to factorize the matrix into two compact low-rank factors, and then optimize these factors directly via simple iterative methods such as gradient descent and alternating minimization. Despite nonconvexity, recent literatures have shown that these simple heuristics in fact achieve linear convergence when initialized properly for a growing number of problems of interest. However, upon closer examination, existing approaches can still be computationally expensive especially for ill-conditioned matrices: the convergence rate of gradient descent depends linearly on the condition number of the low-rank matrix, while the per-iteration cost of alternating minimization is often prohibitive for large matrices. The goal of this paper is to set forth a competitive algorithmic approach dubbed Scaled Gradient Descent (ScaledGD) which can be viewed as preconditioned or diagonally-scaled gradient descent, where the preconditioners are adaptive and iteration-varying with a minimal computational overhead. With tailored variants for low-rank matrix sensing, robust principal component analysis and matrix completion, we theoretically show that ScaledGD achieves the best of both worlds: it converges linearly at a rate independent of the condition number of the low-rank matrix similar as alternating minimization, while maintaining the low per-iteration cost of gradient descent. Our analysis is also applicable to general loss functions that are restricted strongly convex and smooth over low-rank matrices. To the best of our knowledge, ScaledGD is the first algorithm that provably has such properties over a wide range of low-rank matrix estimation tasks. At the core of our analysis is the introduction of a new distance function that takes account of the preconditioners when measuring the distance between the iterates and the ground truth. Finally, numerical examples are provided to demonstrate the effectiveness of ScaledGD in accelerating the convergence rate of ill-conditioned low-rank matrix estimation in a wide number of applications. Tian Tong, Cong Ma 0001, Yuejie Chi |
J. Mach. Learn. Res. | 3 |
| 2021 | Compressed Super-Resolution of Positive SourcesabstractAtomic norm minimization is a convex optimization framework to recover point sources from a subset of their low-pass observations, or equivalently the underlying frequencies of a spectrally-sparse signal. When the amplitudes of the sources are positive, a positive atomic norm can be formulated, and exact recovery can be ensured without imposing a separation between the sources, as long as the number of observations is greater than the number of sources. However, the classic formulation of the atomic norm requires to solve a semidefinite program involving a linear matrix inequality of a size on the order of the signal dimension, which can be prohibitive. In this letter, we introduce a novel “compressed” semidefinite program, which involves a linear matrix inequality of a reduced dimension on the order of the number of sources. We guarantee the tightness of this program under certain conditions on the operator involved in the dimensionality reduction. Finally, we apply the proposed method to direction finding over sparse arrays based on second-order statistics and achieve significant computational savings. Maxime Ferreira Da Costa, Yuejie Chi |
IEEE Signal Process. Lett. | 2 |
| 2021 | Nonconvex Matrix Factorization From Rank-One MeasurementsabstractWe consider the problem of recovering low-rank matrices from random rank-one measurements, which spans numerous applications including covariance sketching, phase retrieval, quantum state tomography, and learning shallow polynomial neural networks, among others. Our approach is to directly estimate the low-rank factor by minimizing a nonconvex least-squares loss function via vanilla gradient descent, following a tailored spectral initialization. When the true rank is bounded by a constant, this algorithm is guaranteed to converge to the ground truth (up to global ambiguity) with near-optimal sample complexity and computational complexity. To the best of our knowledge, this is the first guarantee that achieves near-optimality in both metrics. In particular, the key enabler of near-optimal computational guarantees is an implicit regularization phenomenon: without explicit regularization, both spectral initialization and the gradient descent iterates automatically stay within a region incoherent with the measurement vectors. This feature allows one to employ much more aggressive step sizes compared with the ones suggested in prior literature, without the need of sample splitting. Yuanxin Li 0003, Cong Ma 0001, Yuxin Chen 0002, Yuejie Chi |
IEEE Trans. Inf. Theory | 4 |
| 2021 | Manifold Gradient Descent Solves Multi-Channel Sparse Blind Deconvolution Provably and EfficientlyabstractMulti-channel sparse blind deconvolution, or convolutional sparse coding, refers to the problem of learning an unknown filter by observing its circulant convolutions with multiple input signals that are sparse. This problem finds numerous applications in signal processing, computer vision, and inverse problems. However, it is challenging to learn the filter efficiently due to the bilinear structure of the observations with respect to the unknown filter and inputs, as well as the sparsity constraint. In this paper, we propose a novel approach based on nonconvex optimization over the sphere manifold by minimizing a smooth surrogate of the sparsity-promoting loss function. It is demonstrated that manifold gradient descent with random initializations will probably recover the filter, up to scaling and shift ambiguity, as soon as the number of observations is sufficiently large under an appropriate random data model. Numerical experiments are provided to illustrate the performance of the proposed method with comparisons to existing ones. Laixi Shi, Yuejie Chi |
IEEE Trans. Inf. Theory | 2 |
| 2020 | Communication-Efficient Distributed Optimization in Networks with Gradient Tracking and Variance ReductionabstractDue to the imminent need to alleviate the communication burden in multi-agent and federated learning, the investigation of communication-efficient distributed optimization algorithms for empirical risk minimization has flourished recently. A large fraction of existing algorithms are developed for the master/slave setting, relying on the presence of a central parameter server. This paper focuses on distributed optimization in the network setting (also known as the decentralized setting), where each agent is only allowed to aggregate information from its neighbors over a graph. By properly adjusting the global gradient estimate via a tracking term, we first develop a communication-efficient approximate Newton-type method, called Network-DANE, which generalizes the attractive DANE algorithm to decentralized networks. Our key algorithmic ideas can be applied, in a systematic manner, to obtain decentralized versions of other master/slave distributed algorithms. Notably, we develop Network-SVRG/SARAH, which employ stochastic variance reduction at each agent to accelerate local computations. We establish linear convergence of Network-DANE and Network-SVRG for strongly convex losses, and Network-SARAH for quadratic losses, which shed light on the impact of data homogeneity, network connectivity, and local averaging upon the rate of convergence. Numerical evidence is provided to demonstrate the appealing performance of our algorithms over competitive baselines, in terms of both communication and computation efficiency. Boyue Li, Shicong Cen, Yuxin Chen 0002, Yuejie Chi |
AISTATS | 4 |
| 2020 | Manifold Gradient Descent Solves Multi-Channel Sparse Blind Deconvolution Provably and EfficientlyabstractMulti-channel sparse blind deconvolution refers to the problem of learning an unknown filter by observing its circulant convolutions with multiple input signals that are sparse. It is challenging to learn the filter efficiently due to the bilinear structure of the observations with respect to the unknown filter and inputs, leading to global ambiguities of identification. We propose a novel approach based on nonconvex optimization over the sphere manifold by minimizing a smooth surrogate of the sparsity-promoting loss function. It is demonstrated that manifold gradient descent with random initializations provably recovers the filter, up to scaling and shift ambiguities, as soon as the number of observations is sufficiently large under a suitable random data model. Numerical experiments are conducted to illustrate the efficiency of the proposed method with comparisons to existing methods. Laixi Shi, Yuejie Chi |
ICASSP | 2 |
| 2020 | Sample Complexity of Asynchronous Q-Learning: Sharper Analysis and Variance ReductionabstractAsynchronous Q-learning aims to learn the optimal action-value function (or Q-function) of a Markov decision process (MDP), based on a single trajectory of Markovian samples induced by a behavior policy. Focusing on a $\gamma$-discounted MDP with state space S and action space A, we demonstrate that the $ \ell_{\infty} $-based sample complexity of classical asynchronous Q-learning --- namely, the number of samples needed to yield an entrywise $\epsilon$-accurate estimate of the Q-function --- is at most on the order of $ \frac{1}{ \mu_{\min}(1-\gamma)^5 \epsilon^2 }+ \frac{ t_{\mathsf{mix}} }{ \mu_{\min}(1-\gamma) } $ up to some logarithmic factor, provided that a proper constant learning rate is adopted. Here, $ t_{\mathsf{mix}} $ and $ \mu_{\min} $ denote respectively the mixing time and the minimum state-action occupancy probability of the sample trajectory. The first term of this bound matches the complexity in the case with independent samples drawn from the stationary distribution of the trajectory. The second term reflects the expense taken for the empirical distribution of the Markovian trajectory to reach a steady state, which is incurred at the very beginning and becomes amortized as the algorithm runs. Encouragingly, the above bound improves upon the state-of-the-art result by a factor of at least |S||A|. Further, the scaling on the discount complexity can be improved by means of variance reduction. Gen Li 0005, Yuting Wei 0001, Yuejie Chi, Yuantao Gu, Yuxin Chen 0002 |
NeurIPS | 3 |
| 2020 | Breaking the Sample Size Barrier in Model-Based Reinforcement Learning with a Generative ModelabstractWe investigate the sample efficiency of reinforcement learning in a $\gamma$-discounted infinite-horizon Markov decision process (MDP) with state space S and action space A, assuming access to a generative model. Despite a number of prior work tackling this problem, a complete picture of the trade-offs between sample complexity and statistical accuracy is yet to be determined. In particular, prior results suffer from a sample size barrier, in the sense that their claimed statistical guarantees hold only when the sample size exceeds at least $ |S| |A| / (1-\gamma)^2 $ (up to some log factor). The current paper overcomes this barrier by certifying the minimax optimality of model-based reinforcement learning as soon as the sample size exceeds the order of $ |S| |A| / (1-\gamma) $ (modulo some log factor). More specifically, a perturbed model-based planning algorithm provably finds an $\epsilon$-optimal policy with an order of $ |S| |A| / ((1-\gamma)^3\epsilon^2 ) $ samples (up to log factor) for any $0< \epsilon < 1/(1-\gamma)$. Along the way, we derive improved (instance-dependent) guarantees for model-based policy evaluation. To the best of our knowledge, this work provides the first minimax-optimal guarantee in a generative model that accommodates the entire range of sample sizes (beyond which finding a meaningful policy is information theoretically impossible). Gen Li 0005, Yuting Wei 0001, Yuejie Chi, Yuantao Gu, Yuxin Chen 0002 |
NeurIPS | 3 |
| 2020 | Communication-Efficient Distributed Optimization in Networks with Gradient Tracking and Variance ReductionabstractThere is growing interest in large-scale machine learning and optimization over decentralized networks, e.g. in the context of multi-agent learning and federated learning. Due to the imminent need to alleviate the communication burden, the investigation of communication-efficient distributed optimization algorithms --- particularly for empirical risk minimization --- has flourished in recent years. A large fraction of these algorithms have been developed for the master/slave setting, relying on the presence of a central parameter server that can communicate with all agents. This paper focuses on distributed optimization over networks, or decentralized optimization, where each agent is only allowed to aggregate information from its neighbors over a network (namely, no centralized coordination is present). By properly adjusting the global gradient estimate via local averaging in conjunction with proper correction, we develop a communication-efficient approximate Newton-type method, called Network-DANE, which generalizes DANE to accommodate decentralized scenarios. Our key ideas can be applied, in a systematic manner, to obtain decentralized versions of other master/slave distributed algorithms. A notable development is Network-SVRG/SARAH, which employs variance reduction at each agent to further accelerate local computation. We establish linear convergence of Network-DANE and Network-SVRG for strongly convex losses, and Network-SARAH for quadratic losses, which shed light on the impacts of data homogeneity, network connectivity, and local averaging upon the rate of convergence. We further extend Network-DANE to composite optimization by allowing a nonsmooth penalty term. Numerical evidence is provided to demonstrate the appealing performance of our algorithms over competitive baselines, in terms of both communication and computation efficiency. Our work suggests that by performing a judiciously chosen amount of local communication and computation per iteration, the overall efficiency can be substantially improved. Boyue Li, Shicong Cen, Yuxin Chen 0002, Yuejie Chi |
J. Mach. Learn. Res. | 4 |
| 2020 | On the Stable Resolution Limit of Total Variation Regularization for Spike DeconvolutionabstractThe stability of spike deconvolution, which aims at recovering point sources from their convolution with a point spread function (PSF), is known to be related to the separation between those sources. When the observations are noisy, it is critical to ensure support stability, where the deconvolution does not lead to spurious, or oppositely, missing estimates of the point sources. In this paper, we study the resolution limit of stably recovering the support of two closely located point sources using the Beurling-LASSO estimator, which is a convex optimization approach based on total variation regularization. We establish a sufficient separation criterion between the sources, depending only on the PSF, above which the Beurling-LASSO estimator is guaranteed to return a stable estimate of the point sources, with the same number of estimated elements as that of the ground truth. Our result highlights the impact of PSF on the resolution limit in the noisy setting, which was not evident in previous studies of the noiseless setting. Towards the end, we show that the same resolution limit applies to resolving two close-located sources in conjunction of other well-separated sources. Maxime Ferreira Da Costa, Yuejie Chi |
IEEE Trans. Inf. Theory | 2 |
| 2019 | Nonconvex Matrix Factorization from Rank-One MeasurementsabstractWe consider the problem of recovering low-rank matrices from random rank-one measurements, which spans numerous applications including phase retrieval, quantum state tomography, and learning shallow neural networks with quadratic activations, among others. Our approach is to directly estimate the low-rank factor by minimizing a nonconvex least-squares loss function via vanilla gradient descent, following a tailored spectral initialization. When the true rank is small, this algorithm is guaranteed to converge to the ground truth (up to global ambiguity) with near-optimal sample and computational complexities with respect to the problem size. To the best of our knowledge, this is the first theoretical guarantee that achieves near optimality in both metrics. In particular, the key enabler of near-optimal computational guarantees is an implicit regularization phenomenon: without explicit regularization, both spectral initialization and the gradient descent iterates automatically stay within a region incoherent with the measurement vectors. This feature allows one to employ much more aggressive step sizes compared with the ones suggested in prior literature, without the need of sample splitting. Yuanxin Li 0003, Cong Ma 0001, Yuxin Chen 0002, Yuejie Chi |
AISTATS | 4 |
| 2019 | Shift-invariant Subspace Tracking with Missing DataabstractSubspace tracking is an important problem in signal processing that finds applications in wireless communications, video surveillance, and source localization in radar and sonar. In recent years, it is recognized that a low-dimensional subspace can be estimated and tracked reliably even when the data vectors are partially observed with many missing entries, which is greatly desirable when processing high-dimensional and high-rate data to reduce the sampling requirement. This paper is motivated by the observation that the underlying low-dimensional subspace may possess additional structural properties induced by the physical model of data, which if harnessed properly, can greatly improve subspace tracking performance. As a case study, this paper investigates the problem of tracking direction-of-arrivals from subsampled observations in a unitary linear array, where the signals lie in a subspace spanned by columns of a Vandermonde matrix. We exploit the shift-invariant structure by mapping the data vector to a latent Hankel matrix, and then perform tracking over the Hankel matrices by exploiting their low-rank properties. Numerical simulations are conducted to validate the superiority of the proposed approach over existing subspace tracking methods that do not exploit the additional shift-invariant structure in terms of tracking speed and agility. Myung Cho, Yuejie Chi |
ICASSP | 2 |
| 2019 | On the Sensitivity of Spectral Initialization for Noisy Phase RetrievalabstractThe spectral method is an important approach for signal estimation that is often used as an initialization to iterative methods as well as a stand-alone estimator, where the signal is estimated by the top eigenvector of certain carefully-constructed data matrix. A recent line of work has characterized the asymptotic behavior of such data matrices used in spectral methods, which reveals an interesting phase transition phenomenon: there exists a critical sampling threshold below which the estimate of the spectral method is uninformative. Furthermore, optimal preprocessing functions are developed to minimize this critical sampling threshold. In particular, most of the existing work is focused on the noiseless phase retrieval problem. In this paper, our goal is to examine the sensitivity of such optimal preprocessing functions in noisy phase retrieval, when there is a mismatch between the noise model used in deriving the optimal preprocessing function and the actual noise model in practice. Our results provide important insights into the choice of preprocessing functions in spectral methods. Vincent Monardo, Yuejie Chi |
ICASSP | 2 |
| 2019 | Solving Quadratic Equations via Amplitude-based Nonconvex OptimizationabstractIn many signal processing tasks, one seeks to recover an rcolumn matrix object X ϵ ℂn×rfrom a set of nonnegative quadratic measurements up to orthonormal transforms. Example applications include coherence retrieval in optical imaging and covariance sketching for high-dimensional streaming data. To this end, efficient nonconvex optimization methods are quite appealing, due to their computational efficiency and scalability to large-scale problems. There is a recent surge of activities in designing nonconvex methods for the special case r = 1, known as phase retrieval; however, very little work has studied the general rank-r setting. Motivated by the success of phase retrieval, in this paper we derive several algorithms which utilize the quadratic loss function based on amplitude measurements, including (stochastic) gradient descent and alternating minimization. Numerical experiments demonstrate their computational and statistical performances, highlighting the superior performance of stochastic gradient descent with appropriate mini-batch sizes. Vincent Monardo, Yuanxin Li 0003, Yuejie Chi |
ICASSP | 3 |
| 2019 | Improving Graph Trend Filtering with Non-convex PenaltiesabstractIn this paper, we study the denoising of piecewise smooth graph signals that exhibit inhomogeneous levels of smoothness over a graph. We extend the graph trend filtering framework to a family of nonconvex regularizers that exhibit superior recovery performance over existing convex ones. We present theoretical results in the form of asymptotic error rates for both generic and specialized graph models. We further present an ADMM-based algorithm to solve the proposed optimization problem and analyze its convergence. Numerical performance of the proposed framework with non-convex regularizers on both synthetic and real-world data are presented for denoising, support recovery, and semi-supervised classification. Rohan Varma, Harlin Lee, Yuejie Chi, Jelena Kovacevic |
ICASSP | 3 |
| 2019 | Local Geometry of Cross Entropy Loss in Learning One-Hidden-Layer Neural NetworksabstractWe study model recovery for data classification, where the training labels are generated from a one-hidden-layer neural network with sigmoid activations, and the goal is to recover the weights of the neural network. We consider two network models, the fully-connected network (FCN) and the non-overlapping convolutional neural network (CNN). We prove that with Gaussian inputs, the empirical risk based on cross entropy exhibits strong convexity and smoothness uniformly in a local neighborhood of the ground truth, as soon as the sample complexity is sufficiently large. Hence, if initialized in this neighborhood, it establishes the local convergence guarantee for empirical risk minimization using cross entropy via gradient descent for learning one-hidden-layer neural networks, at the near-optimal sample and computational complexity with respect to the network input dimension without unrealistic assumptions such as requiring a fresh set of samples at each iteration. Haoyu Fu, Yuejie Chi, Yingbin Liang |
ISIT | 2 |
| 2019 | Low-Rank Structured Covariance Matrix EstimationabstractThe covariance matrix estimation problem is posed in both the Bayesian and frequentist settings as the solution of a maximum a posteriori (MAP) or maximum likelihood (ML) optimization, respectively, when the true covariance consists of a known (or bounded) noise floor and a low-rank component. Persymmetric structure may also be assumed. The MAP and ML solutions with the non-convex rank constraint are shown to be a simple scalar thresholding of eigenvalues of a suitably translated and projected sample covariance matrix. No iterative optimization is required; therefore, the computation is suited to real-time applications. Our proof is short and elementary without resorting to the duality theory. Numerical results are presented to illustrate the improved estimation performance obtained by incorporating the structural constraints on the unknown covariance. Azer P. Shikhaliev, Lee C. Potter, Yuejie Chi |
IEEE Signal Process. Lett. | 3 |
| 2018 | Terahertz Imaging of Binary Reflectance with Variational Bayesian InferenceabstractIn this paper, we propose a Bayesian inference approach to extract the binary reflectance pattern of samples from compressed measurements in the terahertz (THz) frequency band. Compared with existing compressed THz imaging methods relying on the sparsity of the reflectance pattern, the proposed Bayesian approach exploits the non-negative binary nature of the reflectance without any assumption on its spatial pattern information and enables a pixel-wise iterative inference approach for fast signal recovery. Numerical evaluation confirms the effectiveness of the proposed approach. Pu Wang 0004, Toshiaki Koike-Akino, Philip V. Orlik, Haoyu Fu, Yuejie Chi |
ICASSP | 5 |
| 2018 | Implicit Regularization in Nonconvex Statistical Estimation: Gradient Descent Converges Linearly for Phase Retrieval and Matrix CompletionabstractRecent years have seen a flurry of activities in designing provably efficient nonconvex optimization procedures for solving statistical estimation problems. For various problems like phase retrieval or low-rank matrix completion, state-of-the-art nonconvex procedures require proper regularization (e.g. trimming, regularized cost, projection) in order to guarantee fast convergence. When it comes to vanilla procedures such as gradient descent, however, prior theory either recommends highly conservative learning rates to avoid overshooting, or completely lacks performance guarantees. This paper uncovers a striking phenomenon in several nonconvex problems: even in the absence of explicit regularization, gradient descent follows a trajectory staying within a basin that enjoys nice geometry, consisting of points incoherent with the sampling mechanism. This “implicit regularization” feature allows gradient descent to proceed in a far more aggressive fashion without overshooting, which in turn results in substantial computational savings. Focusing on two statistical estimation problems, i.e. solving random quadratic systems of equations and low-rank matrix completion, we establish that gradient descent achieves near-optimal statistical and computational guarantees without explicit regularization. As a byproduct, for noisy matrix completion, we demonstrate that gradient descent enables optimal control of both entrywise and spectral-norm errors. Cong Ma 0001, Yuejie Chi, Yuxin Chen 0002 |
ICML | 3 |
| 2018 | Streaming PCA and Subspace Tracking: The Missing Data CaseabstractFor many modern applications in science and engineering, data are collected in a streaming fashion carrying time-varying information, and practitioners need to process them with a limited amount of memory and computational resources in a timely manner for decision making. This often is coupled with the missing data problem, such that only a small fraction of data attributes are observed. These complications impose significant, and unconventional, constraints on the problem of streaming Principal Component Analysis (PCA) and subspace tracking, which is an essential building block for many inference tasks in signal processing and machine learning. This survey article reviews a variety of classical and recent algorithms for solving this problem with low computational and memory complexities, particularly those applicable in the big data regime with missing data. We illustrate that streaming PCA and subspace tracking algorithms can be understood through algebraic and geometric perspectives, and they need to be adjusted carefully to handle missing data. Both asymptotic and non-asymptotic convergence guarantees are reviewed. Finally, we benchmark the performance of several competitive algorithms in the presence of missing data for both well-conditioned and ill-conditioned systems. Laura Balzano, Yuejie Chi, Yue M. Lu |
Proc. IEEE | 2 |
| 2018 | Rethinking PCA for Modern Data Sets: Theory, Algorithms, and ApplicationsabstractThe papers in this special issue introduce the reader to the theory, algorithms, and applications of principal component analysis (PCA) and its many extensions. The aim of PCA is to reduce the dimensionality of multivariate data while preserving as much of the relevant information as possible. It is often the first step in various types of exploratory data analysis, predictive modeling, and classification and clustering tasks, and finds applications in biomedical imaging, computer vision, process fault detection, recommendation systems’ design, and many more domains. Namrata Vaswani, Yuejie Chi, Thierry Bouwmans |
Proc. IEEE | 2 |
| 2018 | Median-Truncated Nonconvex Approach for Phase Retrieval With OutliersabstractThis paper investigates the phase retrieval problem, which aims to recover a signal from the magnitudes of its linear measurements. We develop statistically and computationally efficient algorithms for the situation when the measurements are corrupted by sparse outliers that can take arbitrary values. We propose a novel approach to robustify the gradient descent algorithm by using the sample median as a guide for pruning spurious samples in initialization and local search. Adopting a Poisson loss and a reshaped quadratic loss, respectively, we obtain two algorithms termedmedian-truncated Wirtinger flowandmedian-reshaped Wirtinger flow, both of which provably recover the signal from a near-optimal number of measurements when the measurement vectors are composed of independent and identically distributed Gaussian entries, up to a logarithmic factor, even when a constant fraction of the measurements is adversarially corrupted. We further show that both algorithms are stable in the presence of additional dense bounded noise. Our analysis is accomplished by developing non-trivial concentration results of median-related quantities, which may be of independent interest. We provide numerical experiments to demonstrate the effectiveness of our approach. Huishuai Zhang, Yuejie Chi, Yingbin Liang |
IEEE Trans. Inf. Theory | 2 |
| 2017 | A Nonconvex Approach for Phase Retrieval: Reshaped Wirtinger Flow and Incremental AlgorithmsabstractWe study the problem of solving a quadratic system of equations, i.e., recovering a vector signal $\boldsymbol{x}\in \mathbb{R}^n$ from its magnitude measurements $y_i=|\langle \boldsymbol{a}_i, \boldsymbol{x}\rangle|, i=1,..., m$. We develop a gradient descent algorithm (referred to as RWF for reshaped Wirtinger flow) by minimizing the quadratic loss of the magnitude measurements. Comparing with Wirtinger flow (WF) (Candes et al., 2015), the loss function of RWF is nonconvex and nonsmooth, but better resembles the least-squares loss when the phase information is also available. We show that for random Gaussian measurements, RWF enjoys linear convergence to the true signal as long as the number of measurements is $\mathcal{O}(n)$. This improves the sample complexity of WF ($\mathcal{O}(n\log n)$), and achieves the same sample complexity as truncated Wirtinger flow (TWF) (Chen and Candes, 2015), but without any sophisticated truncation in the gradient loop. Furthermore, RWF costs less computationally than WF, and runs faster numerically than both WF and TWF. We further develop an incremental (stochastic) version of RWF (IRWF) and connect it with the randomized Kaczmarz method for phase retrieval. We demonstrate that IRWF outperforms existing incremental as well as batch algorithms with experiments. Huishuai Zhang, Yingbin Liang, Yuejie Chi |
J. Mach. Learn. Res. | 3 |
| 2016 | Robust blind spikes deconvolutionabstractBlind spikes deconvolution, or blind super-resolution, deals with the problem of estimating the delays and amplitudes of spikes from its convolution with an unknown low-pass point spread function. By constraining the point spread function in a known low-dimensional subspace, a convex optimization algorithm called AtomicLift has been proposed to exactly recover the spikes up to an unavoidable scaling ambiguity in the noiseless setting. This paper analyzes the performance of AtomicLift in the presence of bounded noise, and shows that the spikes are localized in a stable manner where the localization inaccuracy is proportional to the noise level. Moreover, we show AtomicLift is also capable to handle sparse outliers in the frequency domain. Yuejie Chi |
ICASSP | 1 |
| 2016 | Outlier-robust recovery of low-rank positive semidefinite matrices from magnitude measurementsabstractWe address the problem of estimating a low-rank positive semidefinite (PSD) matrix from a set of magnitude measurements that are quadratic in the sensing vectors in the presence of arbitrary outliers. We propose a parameter-free algorithm that seeks the PSD matrix that minimizes the ℓ1-norm of the measurement residual. It is shown that the algorithm can exactly recover a rank-r PSD matrix of size-n from O (nr2) measurements with high probability, even when a fraction of the measurements is corrupted by arbitrary outliers. Furthermore, the recovery is also robust to bounded noise. When an upper bound of the rank of the PSD matrix is known a priori, we further propose a non-convex algorithm based on subgradient descent that demonstrates superior empirical performance. Yuanxin Li 0003, Yuejie Chi |
ICASSP | 3 |
| 2016 | Provable Non-convex Phase Retrieval with Outliers: Median TruncatedWirtinger FlowabstractSolving systems of quadratic equations is a central problem in machine learning and signal processing. One important example is phase retrieval, which aims to recover a signal from only magnitudes of its linear measurements. This paper focuses on the situation when the measurements are corrupted by arbitrary outliers, for which the recently developed non-convex gradient descent Wirtinger flow (WF) and truncated Wirtinger flow (TWF) algorithms likely fail. We develop a novel median-TWF algorithm that exploits robustness of sample median to resist arbitrary outliers in the initialization and the gradient update in each iteration. We show that such a non-convex algorithm provably recovers the signal from a near-optimal number of measurements composed of i.i.d. Gaussian entries, up to a logarithmic factor, even when a constant portion of the measurements are corrupted by arbitrary outliers. We further show that median-TWF is also robust when measurements are corrupted by both arbitrary outliers and bounded noise. Our analysis of performance guarantee is accomplished by development of non-trivial concentration measures of median-related quantities, which may be of independent interest. We further provide numerical experiments to demonstrate the effectiveness of the approach. Huishuai Zhang, Yuejie Chi, Yingbin Liang |
ICML | 2 |
| 2016 | Kaczmarz Method for Solving Quadratic EquationsabstractEstimating low-rank positive-semidefinite (PSD) matrices from symmetric rank-one measurements is of great importance in many applications, such as high-dimensional data processing, quantum state tomography, and phase retrieval. When the rank is known a priori, this problem can be regarded as solving a system of quadratic equations of a low-dimensional subspace. The authors develop a fast iterative algorithm based on an adaptation of the Kaczmarz method, which is traditionally used for solving overdetermined linear systems. In particular, the authors characterize the dynamics of the algorithm when the measurement vectors are composed of standard Gaussian entries in the online setting. Numerical simulations demonstrate the compelling performance of the proposed algorithm. Yuejie Chi, Yue M. Lu |
IEEE Signal Process. Lett. | 1 |
| 2016 | Blind Deconvolution From Multiple Sparse InputsabstractBlind deconvolution is an inverse problem when both the input signal and the convolution kernel are unknown. We propose a convex algorithm based on $1-minimization to solve the blind deconvolution problem, given multiple observations from sparse input signals. The proposed method is related to other problems such as blind calibration and finding sparse vectors in a subspace. Sufficient conditions for exact and stable recovery using the proposed method are developed that shed light on the sample complexity. Finally, numerical examples are provided to showcase the performance of the proposed method. Liming Wang 0004, Yuejie Chi |
IEEE Signal Process. Lett. | 2 |
| 2015 | Compressive graph clustering from random sketchesabstractGraph clustering, where the goal is to cluster the nodes in a graph into disjoint clusters, arises from applications such as community detection, network monitoring, and bioinformatics. This paper describes an approach for graph clustering based on a small number of linear measurements, i.e. sketches, of the adjacency matrix, where each sketch corresponds to the number of edges in a randomly selected subgraph. Under the stochastic block model, we propose a computationally tractable algorithm based on semidefinite programming to recover the underlying clustering structure, by motivating the low-dimensional parsimonious structure of the clustering matrix. Numerical examples are presented to validate the excellent performance of the proposed algorithm, which allows exact recovery of the clustering matrix under favorable trade-offs between the number of sketches and the edge density gap under the stochastic block model. Yuejie Chi |
ICASSP | 1 |
| 2015 | Covariance tracking from sketches of rapid data streamsabstractEstimating and tracking the covariance matrix of high-dimensional data streams with low complexities in acquisition, storage and computation are of great interest in modern data-intensive applications. This paper develops an online covariance estimation and tracking algorithm for a recently developed covariance sketching framework that requires a single sketch per sample [1], by leveraging the low-rank structure of the covariance matrix. In particular, we devise a discounting mechanism in the aggregation procedure to enable faster tracking when the covariance structure changes over time. The performance of the proposed algorithm is validated through numerical examples. Yiran Jiang, Yuejie Chi |
ICASSP | 2 |
| 2015 | Super-resolution of mutually interfering signalsabstractWe consider simultaneously identifying the membership and locations of point sources that are convolved with different low-pass point spread functions, from the observation of their superpositions. This problem arises in three-dimensional super-resolution single-molecule imaging, neural spike sorting, multi-user channel identification, among others. We propose a novel algorithm, based on convex programming, and establish its near-optimal performance guarantee for exact recovery by exploiting the sparsity of the point source model as well as incoherence between the point spread functions. Numerical examples are provided to demonstrate the effectiveness of the proposed approach. Yuanxin Li 0003, Yuejie Chi |
ISIT | 2 |
| 2015 | Orthogonal Matching Pursuit on Faulty CircuitsabstractWith the wide recognition that modern nanoscale devices will be error-prone, characterization of reliability of information processing systems built out of unreliable components has become an important topic. In this paper, we analyze the performance of orthogonal matching pursuit (OMP), a popular sparse recovery algorithm, running on faulty circuits. We identify sufficient conditions for correct recovery of the signal support and express these conditions in terms of the relationship among signal magnitudes, sparsity, and the mutual incoherence of the measurement matrix. We study both the effects of additive errors in arithmetic computations and logical errors in comparators. We find that the additive errors in the OMP computations have an impact on the overall performance comparable to that of the additive noise in the input measurements. We also show that parallel structures are more robust to logical errors than serial structures in the implementation of a noisy arg max operation, and thus lead to a better OMP performance. Yao Li 0007, Yuejie Chi, Chu-Hsiang Huang, Lara Dolecek |
IEEE Trans. Commun. | 2 |
| 2015 | Exact and Stable Covariance Estimation From Quadratic Sampling via Convex ProgrammingabstractStatistical inference and information processing of high-dimensional data often require an efficient and accurate estimation of their second-order statistics. With rapidly changing data, limited processing power and storage at the acquisition devices, it is desirable to extract the covariance structure from a single pass over the data and a small number of stored measurements. In this paper, we explore a quadratic (or rank-one) measurement model which imposes minimal memory requirements and low computational complexity during the sampling process, and is shown to be optimal in preserving various low-dimensional covariance structures. Specifically, four popular structural assumptions of covariance matrices, namely, low rank, Toeplitz low rank, sparsity, jointly rank-one and sparse structure, are investigated, while recovery is achieved via convex relaxation paradigms for the respective structure. The proposed quadratic sampling framework has a variety of potential applications, including streaming data processing, high-frequency wireless communication, phase space tomography and phase retrieval in optics, and noncoherent subspace detection. Our method admits universally accurate covariance estimation in the absence of noise, as soon as the number of measurements exceeds the information theoretic limits. We also demonstrate the robustness of this approach against noise and imperfect structural assumptions. Our analysis is established upon a novel notion called the mixed-norm restricted isometry property (RIP-ℓ2/ℓ1), as well as the conventional RIP-ℓ2/ℓ2for near-isotropic and bounded measurements. In addition, our results improve upon the best-known phase retrieval (including both dense and sparse signals) guarantees using PhaseLift with a significantly simpler approach. Yuxin Chen 0002, Yuejie Chi, Andrea J. Goldsmith |
IEEE Trans. Inf. Theory | 2 |
| 2014 | Estimation of simultaneously structured covariance matrices from quadratic measurementsabstractThis paper explores covariance estimation from energy measurements that are collected via a quadratic form of measurement vectors. A popular structural model is considered where the covariance matrices possess low-rank and sparse structures simultaneously. We investigate a weighted convex relaxation algorithm tailored for this joint structure, which guarantees exact and universal recovery from a small number of measurements. The algorithm is also robust against noise and imperfect structural assumptions. In particular, when the non-zero entries of the covariance matrix exhibit power-law decay, our algorithm admits exact recovery as soon as the number of measurements exceeds the theoretic limit. Our method is related to sparse phase retrieval: the analysis framework herein recovers and strengthens the best-known performance guarantees by extending them to approximately sparse and noisy scenarios as well as a broader class of measurement vectors, and our results are derived using much simpler analysis methods. Yuxin Chen 0002, Yuejie Chi, Andrea J. Goldsmith |
ICASSP | 2 |
| 2014 | Joint sparsity recovery for spectral compressed sensingabstractCompressed Sensing (CS) is an effective approach to reduce the required number of samples for reconstructing a sparse signal in an a priori basis, but may suffer severely from the issue of basis mismatch. In this paper we study the problem of simultaneously recovering multiple spectrally-sparse signals that are supported on the same frequencies lying arbitrarily on the unit circle. We propose an atomic norm minimization problem, which can be regarded as a continuous counterpart of the discrete CS formulation and be solved efficiently via semidefinite programming. Through numerical experiments, we show that the number of samples per signal may be further reduced by harnessing the joint sparsity pattern of multiple signals. Yuejie Chi |
ICASSP | 1 |
| 2014 | Robust and universal covariance estimation from quadratic measurements via convex programmingabstractThis paper considers the problem of recovering the covariance matrix of a stream of high-dimensional data instances from a minimal number of stored measurements. We develop a quadratic random sampling method based on rank-one measurements of the covariance matrix, which serves as an efficient covariance sketching scheme for processing data streams. This also allows modeling of phaseless measurements that arise in high-frequency wireless communication and signal processing applications. We propose to recover the covariance matrix from the above quadratic measurements via convex relaxation with respect to the presumed parsimonious covariance structure. We show that in the absence of noise, exact and universal recovery of low-rank or Toeplitz low-rank covariance matrices can be achieved as soon as the number of stored measurements exceeds the fundamental sampling limit. The convex programs are also robust to noise and imperfect structural assumptions. Our analysis is established upon a novel notion called the mixed-norm restricted isometry property (RIP-ℓ2/ℓ1), as well as the conventional RIP-ℓ2/ℓ2for near-isotropic and bounded measurements. Our results improve upon best-known phase retrieval performance guarantees with a significantly simpler approach. Numerical results are provided to demonstrate the practical applicability of our technique. Yuxin Chen 0002, Yuejie Chi, Andrea J. Goldsmith |
ISIT | 2 |
| 2014 | Classification and Boosting with Multiple Collaborative RepresentationsabstractRecent advances have shown a great potential to explore collaborative representations of test samples in a dictionary composed of training samples from all classes in multi-class recognition including sparse representations. In this paper, we present two multi-class classification algorithms that make use of multiple collaborative representations in their formulations, and demonstrate performance gain of exploring this extra degree of freedom. We first present the Collaborative Representation Optimized Classifier (CROC), which strikes a balance between the nearest-subspace classifier, which assigns a test sample to the class that minimizes the distance between the sample and its principal projection in the selected class, and a Collaborative Representation based Classifier (CRC), which assigns a test sample to the class that minimizes the distance between the sample and its collaborative components. Several well-known classifiers become special cases of CROC under different regularization parameters. We show classification performance can be improved by optimally tuning the regularization parameter through cross validation. We then propose the Collaborative Representation based Boosting (CRBoosting) algorithm, which generalizes the CROC to incorporate multiple collaborative representations. Extensive numerical examples are provided with performance comparisons of different choices of collaborative representations, in particular when the test sample is available via compressive measurements. Yuejie Chi, Fatih Porikli |
IEEE Trans. Pattern Anal. Mach. Intell. | 1 |
| 2014 | Robust Spectral Compressed Sensing via Structured Matrix CompletionabstractThis paper explores the problem of spectral compressed sensing, which aims to recover a spectrally sparse signal from a small random subset of its n time domain samples. The signal of interest is assumed to be a superposition of r multidimensional complex sinusoids, while the underlying frequencies can assume any continuous values in the normalized frequency domain. Conventional compressed sensing paradigms suffer from the basis mismatch issue when imposing a discrete dictionary on the Fourier representation. To address this issue, we develop a novel algorithm, called enhanced matrix completion (EMaC), based on structured matrix completion that does not require prior knowledge of the model order. The algorithm starts by arranging the data into a low-rank enhanced form exhibiting multifold Hankel structure, and then attempts recovery via nuclear norm minimization. Under mild incoherence conditions, EMaC allows perfect recovery as soon as the number of samples exceeds the order of r log4n, and is stable against bounded noise. Even if a constant portion of samples are corrupted with arbitrary magnitude, EMaC still allows exact recovery, provided that the sample complexity exceeds the order of r2log3n. Along the way, our results demonstrate the power of convex relaxation in completing a low-rank multifold Hankel or Toeplitz matrix from minimal observed entries. The performance of our algorithm and its applicability to super resolution are further validated by numerical experiments. Yuxin Chen 0002, Yuejie Chi |
IEEE Trans. Inf. Theory | 2 |
| 2013 | Knowledge-enhanced Matching PursuitabstractCompressive Sensing is possible when the sensing matrix acts as a near isometry on signals of interest that can be sparsely or compressively represented. The attraction of greedy algorithms such as Orthogonal Matching Pursuit is their simplicity. However they fail to take advantage of both the structure of the sensing matrix and any prior information about the sparse signal. This paper introduces an oblique projector to matching pursuit algorithms to enhance detection of a component that is present in the signal by reducing interference from other candidate components based on prior information about the signal as well as the structure of the sensing matrix. Numerical examples demonstrate that performance as a function of SNR is superior to conventional matching pursuit. Yuejie Chi, A. Robert Calderbank |
ICASSP | 1 |
| 2013 | Analysis of fisher information and the Cramer-Rao bound for nonlinear parameter estimation after compressed sensingabstractIn this paper, we analyze the impact of compressed sensing with random matrices on Fisher information and the CRB for estimating unknown parameters in the mean value function of a multivariate normal distribution. We consider the class of random compression matrices that satisfy a version of the Johnson-Lindenstrauss lemma, and we derive analytical lower and upper bounds on the CRB for estimating parameters from randomly compressed data. These bounds quantify the potential loss in CRB as a function of Fisher information of the non-compressed data. In our numerical examples, we consider a direction of arrival estimation problem and compare the actual loss in CRB with our bounds. Pooria Pakrooh, Louis L. Scharf, Ali Pezeshki, Yuejie Chi |
ICASSP | 4 |
| 2013 | Spectral Compressed Sensing via Structured Matrix CompletionabstractThe paper studies the problem of recovering a spectrally sparse object from a small number of time domain samples. Specifically, the object of interest with ambient dimension n is assumed to be a mixture of r complex multi-dimensional sinusoids, while the underlying frequencies can assume any value in the unit disk. Conventional compressed sensing paradigms suffer from the \em basis mismatch issue when imposing a discrete dictionary on the Fourier representation. To address this problem, we develop a novel nonparametric algorithm, called enhanced matrix completion (EMaC), based on structured matrix completion. The algorithm starts by converting the data into a low-rank enhanced form with multi-fold Hankel structure, then attempts recovery via nuclear norm minimization. Under mild incoherence conditions, EMaC allows perfect recovery as soon as the number of samples exceeds the order of \mathcalO(r\log^2 n). We also show that, in many instances, accurate completion of a low-rank multi-fold Hankel matrix is possible when the number of observed entries is proportional to the information theoretical limits (except for a logarithmic gap). The robustness of EMaC against bounded noise and its applicability to super resolution are further demonstrated by numerical experiments. Yuxin Chen 0002, Yuejie Chi |
ICML (3) | 2 |
| 2012 | Connecting the dots in multi-class classification: From nearest subspace to collaborative representationabstractWe present a novel multi-class classifier that strikes a balance between the nearest-subspace classifier, which assigns a test sample to the class that minimizes the distance between the test sample and its principal projection in the selected class, and a collaborative representation based classifier, which classifies a sample to the class that minimizes the distance between the collaborative components of the test sample by using all training samples from all classes as the dictionary and its projection in the selected class. In our formulation, the sparse representation based classifier [1] and nearest subspace classifier become special cases under different regularization parameters. We show that the classification performance can be improved by optimally tuning the regularization parameter, which can be done at almost no extra computational cost. We give extensive numerical examples for digit identification and face recognition with performance comparisons of different choices of collaborative representations, in particular when only a partial observation of the test sample is available via compressive sensing measurements. Yuejie Chi, Fatih Porikli |
CVPR | 1 |
| 2012 | PETRELS: Subspace estimation and tracking from partial observationsabstractWe consider the problem of reconstructing a data stream from a small subset of its entries, where the data stream is assumed to lie in a low-dimensional linear subspace, possibly corrupted by noise. It is also important to track the change of underlying subspace for many applications. This problem can be viewed as a sequential low-rank matrix completion problem in which the subspace is learned in an online fashion. The proposed algorithm, called Parallel Estimation and Tracking by REcursive Least Squares (PETRELS), identifies the underlying low-dimensional subspace via a recursive procedure for each row of the subspace matrix in parallel, and then reconstructs the missing entries via least-squares estimation if required. PETRELS outperforms previous approaches by discounting observations in order to capture long-term behavior of the data stream and be able to adapt to it. Numerical examples are provided for direction-of-arrival estimation and matrix completion, comparing PETRELS with state of the art batch algorithms. Yuejie Chi, Yonina C. Eldar, A. Robert Calderbank |
ICASSP | 1 |
| 2011 | On Training Signal Design for Multi-User MIMO-OFDM: Performance Analysis and TradeoffsabstractThis paper addresses spectrally-efficient multiantenna multi-carrier uplink transmission scenarios where the users overlap in time and frequency and are separated using spatial processing at the base station. The robustness of the proposed training sequences to residual carrier frequency offset and phase noise is evaluated analytically. This analysis reveals an interesting design tradeoff between the Peak-to-Average Power Ratio of a training sequence and the increase in channel estimation mean squared error over the ideal case when these two impairments are not present. Ahmad Gomaa, Yuejie Chi, Naofal Al-Dhahir, A. Robert Calderbank |
VTC Fall | 2 |
| 2011 | Training Signal Design and Tradeoffs for Spectrally-Efficient Multi-User MIMO-OFDM SystemsabstractIn this paper, we design MMSE-optimal training sequences for multi-user MIMO-OFDM systems with an arbitrary number of transmit antennas and an arbitrary number of training symbols. It addresses spectrally-efficient uplink transmission scenarios where the users overlap in time and frequency and are separated using spatial processing at the base station. The robustness of the proposed training sequences to residual carrier frequency offset and phase noise is evaluated. This analysis reveals an interesting design tradeoff between the peak-to-average power ratio of a training sequence and the increase in channel estimation mean squared error over the ideal case when these two impairments are not present. Yuejie Chi, Ahmad Gomaa, Naofal Al-Dhahir, A. Robert Calderbank |
IEEE Trans. Wirel. Commun. | 1 |
| 2010 | Sensitivity to basis mismatch in compressed sensingabstractCompressed sensing theory suggests that successful inversion of an image of the physical world from its modal parameters can be achieved at measurement dimensions far lower than the image dimension, provided that the image is sparse in an a priori known basis. The assumed basis for sparsity typically corresponds to a gridding of the parameter space, e.g., an DFT grid in spectrum analysis. However, in reality no physical field is sparse in the DFT basis or in an a priori known basis. No matter how finely we grid the parameter space the sources may not lie in the center of the grid cells and there is always mismatch between the assumed and the actual bases for sparsity. In this paper, we study the sensitivity of compressed sensing (basis pursuit to be exact) to mismatch between the assumed and the actual sparsity bases. Our mathematical analysis and numerical examples show that the performance of basis pursuit degrades considerably in the presence of basis mismatch. Yuejie Chi, Ali Pezeshki, Louis L. Scharf, A. Robert Calderbank |
ICASSP | 1 |
| 2010 | Compressive blind source separationabstractThe central goal of compressive sensing is to reconstruct a signal that is sparse or compressible in some basis using very few measurements. However reconstruction is often not the ultimate goal and it is of considerable interest to be able to deduce attributes of the signal from the measurements without explicitly reconstructing the full signal. This paper solves the blind source separation problem not in the high dimensional data domain, but in the low dimensional measurement domain. It develops a Bayesian inference framework that integrates hidden Markov models for sources with compressive measurement. Posterior probabilities are calculated using a Markov Chain Monte Carlo (MCMC) algorithm. Simulation results are provided for one-dimensional signals and for two-dimensional images, where hidden Markov tree models of the wavelet coefficients are considered. The integrated Bayesian framework is shown to outperform standard approaches where the mixtures are separated in the data domain. Yiyue Wu, Yuejie Chi, A. Robert Calderbank |
ICIP | 2 |
| 2010 | Regularized blind detection for MIMO communicationsabstractMultiple-Input Multiple-Output (MIMO) systems improve the throughput and reliability of wireless communications. Perfect Channel State Information (CSI) is needed at the receiver to perform coherent detection and achieve the optimal gain of the system. In fast fading and low SNR regimes, it is hard or impossible to obtain perfect CSI, which leads the receiver to operate without knowledge of the CSI and perform blind detection. In reality CSI may be available to the receiver but this CSI may be insufficient to support coherent detection. In this paper, we fill the gap between coherent and blind detection by considering a more realistic model where the receiver knows the statistics of the channel, that is Channel Distribution Information (CDI). We propose a new detection algorithm, called Regularized Blind Detection (RBD), where coherent and blind detection can be viewed as special cases in our model. The algorithm estimates CDI from any training symbols that are available and maximizes performance given the estimated CDI. Simulations demonstrate significant improvement in performance over blind detection. Our work can be viewed as a systematic exploration of space between coherent and blind detection with a strong Bayesian statistic flavor. Yuejie Chi, Yiyue Wu, A. Robert Calderbank |
ISIT | 1 |