Solving Regularized Exp, Cosh and Sinh Regression Problems

Abstract

In modern machine learning, attention computation is a fundamental task for training large language models such as Transformer, GPT-4 and ChatGPT. In this work, we study exponential regression problem which is inspired by the softmax/exp unit in the attention mechanism in large language models. The standard exponential regression is non-convex. We study the regularization version of exponential regression problem which is a convex problem. We use approximate newton method to solve in input sparsity time. Formally, in this problem, one is given matrix $A\in {\mathbb{R}}^{n\times d}$, $b\in {\mathbb{R}}^n$, $w\in {\mathbb{R}}^n$ and any of functions $\exp ,\cosh$ and $\sinh$ denoted as $f$. The goal is to find the optimal $x$ that minimize $0.5\Vert f(Ax)-b{\Vert }_2^2+0.5\Vert \mathrm{diag}(w)Ax{\Vert }_2^2$. The straightforward method is to use the naive Newton's method. Let $\mathrm{nnz}(A)$ denote the number of non-zeros entries in matrix $A$. Let $\omega$ denote the exponent of matrix multiplication. Currently, $\omega \approx 2.373$. Let $\epsilon$ denote the accuracy error. In this paper, we make use of the input sparsity and purpose an algorithm that use $\log (\Vert x_0-x^{\ast }{\Vert }_2/\epsilon )$ iterations and $\tilde{O}(\mathrm{nnz}(A)+d^{\omega })$ per iteration time to solve the problem.