Arvind V. Mahankali

dblp:270/8242 · DBLP profile ↗
← Back
6ranked-venue papers
5as first author
6since 2021 · last 2024
—ORCID · none

Domains — the database's venue-derived domains; a paper can count in several

Artificial intelligence and machine learning · 4 · 3 first-author · 4 since 2021Theory of computation · 2 · 2 first-author · 2 since 2021
YearPublicationVenuePosition
2024 One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention
abstract
Recent works have empirically analyzed in-context learning and shown that transformers trained on synthetic linear regression tasks can learn to implement ridge regression, which is the Bayes-optimal predictor, given sufficient capacity (Akyurek et al., 2023), while one-layer transformers with linear self-attention and no MLP layer will learn to implement one step of gradient descent (GD) on a least-squares linear regression objective (von Oswald et al., 2022). However, the theory behind these observations remains poorly understood. We theoretically study transformers with a single layer of linear self-attention, trained on synthetic noisy linear regression data. First, we mathematically show that when the covariates are drawn from a standard Gaussian distribution, the one-layer transformer which minimizes the pre-training loss will implement a single step of GD on the least-squares linear regression objective. Then, we find that changing the distribution of the covariates and weight vector to a non-isotropic Gaussian distribution has a strong impact on the learned algorithm: the global minimizer of the pre-training loss now implements a single step of $\textit{pre-conditioned}$ GD. However, if only the distribution of the responses is changed, then this does not have a large effect on the learned algorithm: even when the response comes from a more general family of $\textit{nonlinear}$ functions, the global minimizer of the pre-training loss still implements a single step of GD on a least-squares linear regression objective.
Arvind V. Mahankali, Tatsunori B. Hashimoto, Tengyu Ma 0001
ICLR1
2024 Near-Linear Time and Fixed-Parameter Tractable Algorithms for Tensor Decompositions
abstract
We study low rank approximation of tensors, focusing on the Tensor Train and Tucker decompositions, as well as approximations with tree tensor networks and general tensor networks. As suggested by hardness results also shown in this work, obtaining (1+ε)-approximation algorithms for rank k tensor train and Tucker decompositions efficiently may be computationally hard for these problems. Therefore, we propose different algorithms that respectively satisfy some of the objectives above while violating some others within a bound, known as bicriteria algorithms. On the one hand, for rank-k tensor train decomposition for tensors with q modes, we give a (1 + ε)-approximation algorithm with a small bicriteria rank (O(qk/ε) up to logarithmic factors) and O(q ⋅ nnz(A)) running time, up to lower order terms. Here nnz(A) denotes the number of non-zero entries in the input tensor A. We also show how to convert the algorithm of [Huber et al., 2017] into a relative error approximation algorithm, but their algorithm necessarily has a running time of O(qr² ⋅ nnz(A)) + n ⋅ poly(qk/ε) when converted to a (1 + ε)-approximation algorithm with bicriteria rank r. Thus, the running time of our algorithm is better by at least a k² factor. To the best of our knowledge, our work is the first to achieve a near-input-sparsity time relative error approximation algorithm for tensor train decomposition. Our key technique is a method for efficiently obtaining subspace embeddings for a matrix which is the flattening of a Tensor Train of q tensors - the number of rows in the subspace embeddings is polynomial in q, thus avoiding the curse of dimensionality. We extend our algorithm to tree tensor networks and tensor networks on arbitrary graphs. Another way of coping with intractability is by looking at fixed-parameter tractable (FPT) algorithms. We give FPT algorithms for the tensor train, Tucker, and Canonical Polyadic (CP) decompositions, which are simpler than the FPT algorithms of [Song et al., 2019], since our algorithms do not make use of polynomial system solvers. Our technique of using an exponential number of Gaussian subspace embeddings with exactly k rows (and thus exponentially small success probability) may be of independent interest.
Arvind V. Mahankali, David P. Woodruff, Ziyu Zhang 0002
ITCS1
2023 Beyond NTK with Vanilla Gradient Descent: A Mean-Field Analysis of Neural Networks with Polynomial Width, Samples, and Time
abstract
Despite recent theoretical progress on the non-convex optimization of two-layer neural networks, it is still an open question whether gradient descent on neural networks without unnatural modifications can achieve better sample complexity than kernel methods. This paper provides a clean mean-field analysis of projected gradient flow on polynomial-width two-layer neural networks. Different from prior works, our analysis does not require unnatural modifications of the optimization algorithm. We prove that with sample size $n = O(d^{3.1})$ where $d$ is the dimension of the inputs, the network trained with projected gradient flow converges in polynomial time to a non-trivial error that is not achievable by kernel methods using $n \ll d^4$ samples, hence demonstrating a clear separation between unmodified gradient descent and NTK. As a corollary, we show that projected gradient descent with a positive learning rate and a polynomial number of iterations converges to low error with the same sample complexity.
Arvind V. Mahankali, Kefan Dong, Margalit Glasgow, Tengyu Ma 0001
NeurIPS1
2021 Streaming and Distributed Algorithms for Robust Column Subset Selection
abstract
We give the first single-pass streaming algorithm for Column Subset Selection with respect to the entrywise $\ell_p$-norm with $1 \leq p < 2$. We study the $\ell_p$ norm loss since it is often considered more robust to noise than the standard Frobenius norm. Given an input matrix $A \in \mathbb{R}^{d \times n}$ ($n \gg d$), our algorithm achieves a multiplicative $k^{\frac{1}{p} - \frac{1}{2}}\poly(\log nd)$-approximation to the error with respect to the \textit{best possible column subset} of size $k$. Furthermore, the space complexity of the streaming algorithm is optimal up to a logarithmic factor. Our streaming algorithm also extends naturally to a 1-round distributed protocol with nearly optimal communication cost. A key ingredient in our algorithms is a reduction to column subset selection in the $\ell_{p,2}$-norm, which corresponds to the $p$-norm of the vector of Euclidean norms of each of the columns of $A$. This enables us to leverage strong coreset constructions for the Euclidean norm, which previously had not been applied in this context. We also give the first provable guarantees for greedy column subset selection in the $\ell_{1, 2}$ norm, which can be used as an alternative, practical subroutine in our algorithms. Finally, we show that our algorithms give significant practical advantages on real-world data analysis tasks.
Shuli Jiang, Dennis Li, Irene Mengze Li, Arvind V. Mahankali, David P. Woodruff
ICML4
2021 Linear and Kernel Classification in the Streaming Model: Improved Bounds for Heavy Hitters
abstract
We study linear and kernel classification in the streaming model. For linear classification, we improve upon the algorithm of (Tai, et al. 2018), which solves the $\ell_1$ point query problem on the optimal weight vector $w_* \in \mathbb{R}^d$ in sublinear space. We first give an algorithm solving the more difficult $\ell_2$ point query problem on $w_*$, also in sublinear space. We also give an algorithm which solves the $\ell_2$ heavy hitter problem on $w_*$, in sublinear space and running time. Finally, we give an algorithm which can $\textit{deterministically}$ solve the $\ell_1$ point query problem on $w_*$, with sublinear space improving upon that of (Tai, et al. 2018). For kernel classification, if $w_* \in \mathbb{R}^{d^p}$ is the optimal weight vector classifying points in the stream according to their $p^{th}$-degree polynomial kernel, then we give an algorithm solving the $\ell_2$ point query problem on $w_*$ in $\text{poly}(\frac{p \log d}{\varepsilon})$ space, and an algorithm solving the $\ell_2$ heavy hitter problem in $\text{poly}(\frac{p \log d}{\varepsilon})$ space and running time. Note that our space and running time are polynomial in $p$, making our algorithms well-suited to high-degree polynomial kernels and the Gaussian kernel (approximated by the polynomial kernel of degree $p = \Theta(\log T)$).
Arvind V. Mahankali, David P. Woodruff
NeurIPS1
2021 Optimal ℓ1 Column Subset Selection and a Fast PTAS for Low Rank Approximation
abstract
We study the problem of entrywise ℓ1 low rank approximation. We give the first polynomial time column subset selection-based ℓ1 low rank approximation algorithm sampling Õ(k) columns and achieving an Õ(k1/2)-approximation for any k, improving upon the previous best Õ(k)-approximation and matching a prior lower bound for column subset selection-based ℓ1-low rank approximation which holds for any poly(k) number of columns. We extend our results to obtain tight upper and lower bounds for column subset selection-based ℓp low rank approximation for any 1 < p < 2, closing a long line of work on this problem. We next give a (1 + ∊)-approximation algorithm for entrywise ℓp low rank approximation of an n × d matrix, for 1 ≤ p < 2, that is not a column subset selection algorithm. First, we obtain an algorithm which, given a matrix A ∊ ℝn × d, returns a rank-k matrix  in 2poly(k/∊) + poly(nd) running time that achieves the following guarantee: where . Using this algorithm, in the same running time we give an algorithm which obtains error at most (1 + ∊) · OPT and outputs a matrix of rank at most 3k — these algorithms significantly improve upon all previous (1 + ∊)- and O(1)-approximation algorithms for the ℓp low rank approximation problem, which required at least npoly(k/∊) or npoly(k) running time, and either required strong bit complexity assumptions (our algorithms do not) or had bicriteria rank 3k. Finally, we show hardness results which nearly match our 2poly(k) +poly(nd) running time and the above additive error guarantee.
Arvind V. Mahankali, David P. Woodruff
SODA1