publications
in reverse chronological order.
2026
- arXivDouble Preconditioning (DoPr): Optimization for Test-Time Performance, not Validation LossSubmitted , 2026
Many modern applications of deep learning involve training a neural network via a one-step prediction loss (e.g., L^2 regression, cross-entropy), but deploy the network by rolling out along its own predictions. Key examples include autoregressive language modeling, flow-based generative modeling, and robot policy learning. It is well-documented that these settings induce a phenomenon we call test-time feedback (TTF): the mismatch between the training/validation loss and downstream metrics of interest, such as task success rate and generation quality, which grows with task length. While data curation, architecture, and objective design have been proposed to combat train-test shift in TTF settings, this paper proposes optimization as a new design axis to mitigate error accumulation. Specifically, we introduce a new optimization paradigm called double-preconditioning (DoPr) uniquely tailored to the challenges of TTF. DoPr combines gradient-wise preconditioning, as in Adam and Muon, with activation-wise preconditioning (AP), such as in KFAC. We show that the addition of AP yields a drop-in intervention for increasing downstream model performance across a range of TTF settings. Interestingly, these gains in test-time performance do not consistently accompany improvements in validation loss, opening new questions about how to properly evaluate models trained with one-step supervised objectives.
- ICLRAction Chunking and Exploratory Data Collection Yield Exponential Improvements in Behavior Cloning for Continuous ControlIn International Conference on Learning Representations (ICLR), 2026
This paper presents a theoretical analysis of two of the most impactful interventions in modern learning from demonstration in robotics and continuous control: the practice of action-chunking (predicting sequences of actions in open-loop) and exploratory augmentation of expert demonstrations. Though recent results show that learning from demonstration, also known as imitation learning (IL), can suffer errors that compound exponentially with task horizon in continuous settings, we demonstrate that action chunking and exploratory data collection circumvent exponential compounding errors in different regimes. Our results identify control-theoretic stability as the key mechanism underlying the benefits of these interventions. On the empirical side, we validate our predictions and the role of control-theoretic stability through experimentation on popular robot learning benchmarks. On the theoretical side, we demonstrate that the control-theoretic lens provides fine-grained insights into how compounding error arises, leading to tighter statistical guarantees on imitation learning error when these interventions are applied than previous techniques based on information-theoretic considerations alone.
2025
- ICMLOn The Concurrence of Layer-wise Preconditioning Methods and Provable Feature LearningThomas T. Zhang, Behrad Moniri, Ansh Nagwekar, Faraz Rahman, Anton Xue, Hamed Hassani, and Nikolai MatniIn International Conference on Machine Learning (ICML), 2025
Layer-wise preconditioning methods are a family of memory-efficient optimization algorithms that introduce preconditioners per axis of each layer’s weight tensors. These methods have seen a recent resurgence, demonstrating impressive performance relative to entry-wise ("diagonal") preconditioning methods such as Adam(W) on a wide range of neural network optimization tasks. Complementary to their practical performance, we demonstrate that layer-wise preconditioning methods are provably necessary from a statistical perspective. To showcase this, we consider two prototypical models, linear representation learning and single-index learning, which are widely used to study how typical algorithms efficiently learn useful features to enable generalization. In these problems, we show SGD is a suboptimal feature learner when extending beyond ideal isotropic inputs \mathbf{x} ∼\mathsf{N}(\mathbf{0}, \mathbf{I}) and well-conditioned settings typically assumed in prior work. We demonstrate theoretically and numerically that this suboptimality is fundamental, and that layer-wise preconditioning emerges naturally as the solution. We further show that standard tools like Adam preconditioning and batch-norm only mildly mitigate these issues, supporting the unique benefits of layer-wise preconditioning.
2024
- AAAIRegret Analysis of Multi-task Representation Learning for Linear-Quadratic Adaptive ControlIn AAAI Conference on Artificial Intelligence, 2024
Representation learning is a powerful tool that enables learning over large multitudes of agents or domains by enforcing that all agents operate on a shared set of learned features. However, many robotics or controls applications that would benefit from collaboration operate in settings with changing environments and goals, whereas most guarantees for representation learning are stated for static settings. Toward rigorously establishing the benefit of representation learning in dynamic settings, we analyze the regret of multi-task representation learning for linear-quadratic control. This setting introduces unique challenges. Firstly, we must account for and balance the \textit{misspecification} introduced by an approximate representation. Secondly, we cannot rely on the parameter update schemes of single-task online LQR, for which least-squares often suffices, and must devise a novel scheme to ensure sufficient improvement. We demonstrate that for settings where exploration is "benign", the regret of any agent after T timesteps scales as \tilde O(\sqrt{T/H}), where H is the number of agents. In settings with "difficult" exploration, the regret scales as \tilde O(\sqrt{d_u d_θ} \sqrt{T} + T^{3/4}/H^{1/5}), where d_x is the state-space dimension, d_u is the input dimension, and d_θ is the task-specific parameter count. In both cases, by comparing to the minimax single-task regret O(\sqrt{d_x d_u^2}\sqrt{T}), we see a benefit of a large number of agents. Notably, in the difficult exploration case, by sharing a representation across tasks, the effective task-specific parameter count can often be small d_θ< d_x d_u. Lastly, we provide numerical validation of the trends we predict.
- ICMLGuarantees for Nonlinear Representation Learning: Non-identical Covariates, Dependent Data, Fewer SamplesIn International Conference on Machine Learning (ICML), 2024
A driving force behind the diverse applicability of modern machine learning is the ability to extract meaningful features across many sources. However, many practical domains involve data that are non-identically distributed across sources, and statistically dependent within its source, violating vital assumptions in existing theoretical studies. Toward addressing these issues, we establish statistical guarantees for learning general \textit{nonlinear} representations from multiple data sources that admit different input distributions and possibly dependent data. Specifically, we study the sample-complexity of learning T+1 functions f_⋆^{(t)} ∘g_⋆from a function class \mathcal F \times \mathcal G, where f_⋆^{(t)} are task specific linear functions and g_⋆is a shared nonlinear representation. A representation \hat g is estimated using N samples from each of T source tasks, and a fine-tuning function \hat f^{(0)} is fit using N’ samples from a target task passed through \hat g. We show that when N ≳C_{\mathrm{dep}} (\mathrm{dim}(\mathcal F) + \mathrm{C}(\mathcal G)/T), the excess risk of \hat f^{(0)} ∘\hat g on the target task decays as ν_{\mathrm{div}} \big(\frac{\mathrm{dim}(\mathcal F)}{N’} + \frac{\mathrm{C}(\mathcal G)}{N T} \big), where C_{\mathrm{dep}} denotes the effect of data dependency, ν_{\mathrm{div}} denotes an (estimatable) measure of \textit{task-diversity} between the source and target tasks, and \mathrm C(\mathcal G) denotes the complexity of the representation class \mathcal G. In particular, our analysis reveals: as the number of tasks T increases, both the sample requirement and risk bound converge to that of r-dimensional regression as if g_⋆had been given, and the effect of dependency only enters the sample requirement, leaving the risk bound matching the iid setting.
- ICLRSample-Efficient Linear Representation Learning from Non-IID Non-Isotropic DataIn International Conference on Learning Representations (ICLR), 2024
Spotlight presentation
A powerful concept behind much of the recent progress in machine learning is the extraction of common features across data from heterogeneous sources or tasks. Intuitively, using all of one’s data to learn a common representation function benefits both computational effort and statistical generalization by leaving a smaller number of parameters to fine-tune on a given task. Toward theoretically grounding these merits, we propose a general setting of recovering linear operators M from noisy vector measurements y = Mx + w, where the covariates x may be both non-i.i.d. and non-isotropic. We demonstrate that existing isotropy-agnostic representation learning approaches incur biases on the representation update, which causes the scaling of the noise terms to lose favorable dependence on the number of source tasks. This in turn can cause the sample complexity of representation learning to be bottlenecked by the single-task data size. We introduce an adaptation, \texttt{De-bias & Feature-Whiten} (\texttt{DFW}), of the popular alternating minimization-descent scheme proposed independently in Collins et al., (2021) and Nayer and Vaswani (2022), and establish linear convergence to the optimal representation with noise level scaling down with the \textit{total} source data size. This leads to generalization bounds on the same order as an oracle empirical risk minimizer. We verify the vital importance of \texttt{DFW} on various numerical simulations. In particular, we show that vanilla alternating-minimization descent fails catastrophically even for iid, but mildly non-isotropic data. Our analysis unifies and generalizes prior work, and provides a flexible framework for a wider range of applications, such as in controls and dynamical systems.
2023
- TACPerformance-Robustness Tradeoffs in Adversarially Robust Control and EstimationIEEE Transactions on Automatic Control, 2023
While \mathcal{H}_∞methods can introduce robustness against worst-case perturbations, their nominal performance under conventional stochastic disturbances is often drastically reduced. Though this fundamental tradeoff between nominal performance and robustness is known to exist, it is not well-characterized in quantitative terms. Toward addressing this issue, we borrow the increasingly ubiquitous notion of adversarial training from machine learning to construct a class of controllers which are optimized for disturbances consisting of mixed stochastic and worst-case components. We find that this problem admits a linear time invariant optimal controller that has a form closely related to suboptimal \mathcal{H}_∞solutions. We then provide a quantitative performance-robustness tradeoff analysis in two analytically tractable cases: state feedback control, and state estimation. In these special cases, we demonstrate that the severity of the tradeoff depends in an interpretable manner upon system-theoretic properties such as the spectrum of the controllability gramian, the spectrum of the observability gramian, and the stability of the system. This provides practitioners with general guidance for determining how much robustness to incorporate based on a priori system knowledge. We empirically validate our results by comparing the performance of our controller against standard baselines, and plotting tradeoff curves.
- L4DCMulti-Task Imitation Learning for Linear Dynamical SystemsThomas T. Zhang, Katie Kang, Bruce D. Lee, Claire Tomlin, Sergey Levine, Stephen Tu, and Nikolai MatniIn Learning for Dynamics and Control Conference (L4DC), 2023
We study representation learning for efficient imitation learning over linear systems. In particular, we consider a setting where learning is split into two phases: (a) a pre-training step where a shared k-dimensional representation is learned from H source policies, and (b) a target policy fine-tuning step where the learned representation is used to parameterize the policy class. We find that the imitation gap over trajectories generated by the learned target policy is bounded by \tilde{O}\left( \frac{k n_x}{HN_{\mathrm{shared}}} + \frac{k n_u}{N_{\mathrm{target}}}\right), where n_x > k is the state dimension, n_u is the input dimension, N_{\mathrm{shared}} denotes the total amount of data collected for each policy during representation learning, and N_{\mathrm{target}} is the amount of target task data. This result formalizes the intuition that aggregating data across related tasks to learn a representation can significantly improve the sample efficiency of learning a target task. The trends suggested by this bound are corroborated in simulation.
- ACCAdversarial Tradeoffs in Robust State EstimationIn American Control Conference (ACC), 2023
Adversarially robust training has been shown to reduce the susceptibility of learned models to targeted input data perturbations. However, it has also been observed that such adversarially robust models suffer a degradation in accuracy when applied to unperturbed data sets, leading to a robustness-accuracy tradeoff. Inspired by recent progress in the adversarial machine learning literature which characterize such tradeoffs in simple settings, we develop tools to quantitatively study the performance-robustness tradeoff between nominal and robust state estimation. In particular, we define and analyze a novel \textit{adversarially robust Kalman Filtering problem}. We show that in contrast to most problem instances in adversarial machine learning, we can precisely derive the adversarial perturbation in the Kalman Filtering setting. We provide an algorithm to find this perturbation given data realizations, and develop upper and lower bounds on the adversarial state estimation error in terms of the standard (non-adversarial) estimation error and the spectral properties of the resulting observer. Through these results, we show a natural connection between a filter’s robustness to adversarial perturbation and underlying control theoretic properties of the system being observed, namely the spectral properties of its observability gramian.
2022
- NeurIPSTaSIL: Taylor Series Imitation LearningIn Advances in Neural Information Processing Systems (NeurIPS), 2022
We propose Taylor Series Imitation Learning (TaSIL), a simple augmentation to standard behavior cloning losses in the context of continuous control. TaSIL penalizes deviations in the higher-order Taylor series terms between the learned and expert policies. We show that experts satisfying a notion of \textit{incremental input-to-state stability} are easy to learn, in the sense that a small TaSIL-augmented imitation loss over expert trajectories guarantees a small imitation loss over trajectories generated by the learned policy. We provide sample-complexity bounds for TaSIL that scale as \tilde{\mathcal{O}}(1/n) in the realizable setting, for n the number of expert demonstrations. Finally, we demonstrate experimentally the relationship between the robustness of the expert policy and the order of Taylor expansion required in TaSIL, and compare standard Behavior Cloning, DART, and DAgger with TaSIL-loss-augmented variants. In all cases, we show significant improvement over baselines across a variety of MuJoCo tasks.
- CDCPerformance-Robustness Tradeoffs in Adversarially Robust Linear-Quadratic ControlIn IEEE Conference on Decision and Control (CDC), 2022
While \mathcal{H}_∞methods can introduce robustness against worst-case perturbations, their nominal performance under conventional stochastic disturbances is often drastically reduced. Though this fundamental tradeoff between nominal performance and robustness is known to exist, it is not well-characterized in quantitative terms. Toward addressing this issue, we borrow from the increasingly ubiquitous notion of adversarial training from machine learning to construct a class of controllers which are optimized for disturbances consisting of mixed stochastic and worst-case components. We find that this problem admits a stationary optimal controller that has a simple analytic form closely related to suboptimal \mathcal{H}_∞solutions. We then provide a quantitative performance-robustness tradeoff analysis, in which system-theoretic properties such as controllability and stability explicitly manifest in an interpretable manner. This provides practitioners with general guidance for determining how much robustness to incorporate based on a priori system knowledge. We empirically validate our results by comparing the performance of our controller against standard baselines, and plotting tradeoff curves.
- L4DCAdversarially Robust Stability Certificates can be Sample-EfficientIn Learning for Dynamics and Control Conference (L4DC), 2022
Motivated by bridging the simulation to reality gap in the context of safety-critical systems, we consider learning adversarially robust stability certificates for unknown nonlinear dynamical systems. In line with approaches from robust control, we consider additive and Lipschitz bounded adversaries that perturb the system dynamics. We show that under suitable assumptions of incremental stability on the underlying system, the statistical cost of learning an adversarial stability certificate is equivalent, up to constant factors, to that of learning a nominal stability certificate. Our results hinge on novel bounds for the Rademacher complexity of the resulting adversarial loss class, which may be of independent interest. To the best of our knowledge, this is the first characterization of sample-complexity bounds when performing adversarial learning over data generated by a dynamical system. We further provide a practical algorithm for approximating the adversarial training algorithm, and validate our findings on a damped pendulum example.
2021
- SIMODSBiwhitening Reveals the Rank of a Count MatrixBoris Landa, Thomas T. C. K. Zhang, and Yuval KlugerSIAM Journal on Mathematics of Data Science, 2021
Estimating the rank of a corrupted data matrix is an important task in data analysis, most notably for choosing the number of components in PCA. Significant progress on this task was achieved using random matrix theory by characterizing the spectral properties of large noise matrices. However, utilizing such tools is not straightforward when the data matrix consists of count random variables, e.g., Poisson, in which case the noise can be heteroskedastic with an unknown variance in each entry. In this work, we consider a Poisson random matrix with independent entries, and propose a simple procedure termed \textit{biwhitening} for estimating the rank of the underlying signal matrix (i.e., the Poisson parameter matrix) without any prior knowledge. Our approach is based on the key observation that one can scale the rows and columns of the data matrix simultaneously so that the spectrum of the corresponding noise agrees with the standard Marchenko-Pastur (MP) law, justifying the use of the MP upper edge as a threshold for rank selection. Importantly, the required scaling factors can be estimated directly from the observations by solving a matrix scaling problem via the Sinkhorn-Knopp algorithm. Aside from the Poisson, our approach is extended to families of distributions that satisfy a quadratic relation between the mean and the variance, such as the generalized Poisson, binomial, negative binomial, gamma, and many others. This quadratic relation can also account for missing entries in the data. We conduct numerical experiments that corroborate our theoretical findings, and showcase the advantage of our approach for rank estimation in challenging regimes. Furthermore, we demonstrate the favorable performance of our approach on several real datasets of single-cell RNA sequencing (scRNA-seq), High-Throughput Chromosome Conformation Capture (Hi-C), and document topic modeling.
2019
- arXivSparse Optimization on General Atomic Sets: Greedy and Forward-Backward AlgorithmsThomas ZhangPreprint , 2019
We consider the problem of sparse atomic optimization, where the notion of "sparsity" is generalized to meaning some linear combination of few atoms. The definition of atomic set is very broad; popular examples include the standard basis, low-rank matrices, overcomplete dictionaries, permutation matrices, orthogonal matrices, etc. The model of sparse atomic optimization therefore includes problems coming from many fields, including statistics, signal processing, machine learning, computer vision and so on. Specifically, we consider the problem of maximizing a restricted strongly convex (or concave), smooth function restricted to a sparse linear combination of atoms. We extend recent work that establish linear convergence rates of greedy algorithms on restricted strongly concave, smooth functions on sparse vectors to the realm of general atomic sets, where the convergence rate involves a novel quantity: the "sparse atomic condition number". This leads to the strongest known multiplicative approximation guarantees for various flavors of greedy algorithms for sparse atomic optimization; in particular, we show that in many settings of interest the greedy algorithm can attain strong approximation guarantees while maintaining sparsity. Furthermore, we introduce a scheme for forward-backward algorithms that achieves the same approximation guarantees. Secondly, we define an alternate notion of weak submodularity, which we show is tightly related to the more familiar version that has been used to prove earlier linear convergence rates. We prove analogous multiplicative approximation guarantees using this alternate weak submodularity, and establish its distinct identity and applications.
2014
- arXivOn the continuous Fermat-Weber problem for a convex polygon using Euclidean distanceThomas T. C. K. Zhang and John Gunnar Carlsson2014
We consider the continuous Fermat-Weber problem, where the customers are continuously (uniformly) distributed along the boundary of a convex polygon. We derive the closed-form expression for finding the average distance from a given point to the continuously distributed customers along the boundary. A Weiszfeld-type procedure is proposed for this model, which is shown to be linearly convergent. We also derive a closed-form formula to find the average distance for a given point to the entire convex polygon, assuming a uniform distribution. Since the function is smooth, convex, and explicitly given, the continuous version of the Fermat-Weber problem over a convex polygon can be solved easily by numerical algorithms.