Attention Scheme Inspired Softmax Regression

Introduction

In the past few years, Large Language Models (LLMs) have experienced an explosive development. There is a series of results of LLMs, like Transformer , GPT-1 , BERT , GPT-2 , GPT-3 , PaLM , OPT . The success of a recent chatbot named ChatGPT by OpenAI has exemplified the use of LLMs in human-interaction tasks. Very recently, OpenAI released their new version of LLM, named GPT-4 , which has been tested to perform much better even than previous ChatGPT . These LLMs are trained on massive amounts of textual data to generate natural language text. They have already shown their power on various real-work tasks, including natural language translation , sentiment analysis , language modeling , and even creative writing .

In the construction of the LLMs, attention computation is a key component which is used to enhance the model’s ability to focus on relevant parts of the input text . The attention matrix is defined as a squared matrix consisted of rows and columns related to the words or tokens, and the entries in the matrix represent the correlations between the words/tokens in the input text. The attention mechanism allows the model to selectively focus on specific parts of the input text when generating the output, rather than treating all input tokens equally. The attention mechanism is based on the idea that different parts of the input sequence contribute differently to the output sequence, and the model should learn to weigh these contributions accordingly. In LLMs, attention computation is typically implemented as a soft attention mechanism, where the weights are computed using a softmax function over the input sequence. The Attention computation can be described as follows (see as an example).

Motivated by the exp function in attention computation, previous work has formally defined hyperbolic function (for example $f(x)=\exp(Ax),\cosh(Ax),\sinh(Ax)$) regression problems.

Here $f(x)$ can be either of $\exp(Ax)$, $\cosh(Ax)$ and $\sinh(Ax)$.

In this work, we move one more step forward and to consider the normalization factor, $\langle f(x),{\bf 1}_{n}\rangle^{-1}$ $=\langle\exp(Ax),{\bf 1}_{n}\rangle^{-1}$. We will focus on the exp in the rest of the paper. Inspired by the softmax formulation in each row of the above attention computation, we formally define the softmax regression problem,

It is natural in practice to consider regularization , then we consider the regularized version of softmax regression.

We define $f(x):=\langle\exp(Ax),\mathbf{1}_{n}\rangle^{-1}\exp(Ax)$.

We use $x^{*}$ to denote the optimal solution of

Assume that $\|A\|\leq R$. Here $\|A\|$ denotes the spectral norm of matrix $A$.

Suppose that $b\geq{\bf 0}_{n}$ and $\|b\|_{1}\leq 1$. Here ${\bf 0}_{n}$ denotes a length-$n$ vector where all the entries are zeros.

Assume that $w_{i}^{2}\geq 100+l/\sigma_{\min}(A)^{2}$ for all $i\in[n]$. Here $\sigma_{\min}(A)$ denotes the smallest singular value of matrix $A$.

Let $x_{0}$ denote an starting/initial point such that $M\|x_{0}-x^{*}\|_{2}\leq 0.1l$.

Let $\epsilon\in(0,0.1)$ be our accuracy parameter.

Let $\delta\in(0,0.1)$ be our failure probability.

Let $\omega$ denote the exponent of matrix multiplication.

There is a randomized algorithm (Algorithm 1) that

runs $\log(\|x_{0}-x^{*}\|_{2}/\epsilon)$ iterations

We remark that, in previous work , they only assume $\|b\|_{2}\leq R$. The reason is in their setting, they don’t consider the normalization parameter. It doesn’t make sense for them to assume that $\|b\|_{1}\leq 1$ because they’re not trying to learn the distribution.

We organize the following paper as follows. In Section 2, we introduce some other projects that’s related to or that has inspired our work. In Section 3 we provide a sketch for the techniques used in our project. In Section 4 we define the notations used in our work and provide some useful tools for exact algebra, approximate algebra and differential computation. In Section 5 we provide detailed analysis of $L_{\exp}$, including its gradient and hessian. In Section 6 we proved that $L=L_{\exp}+L_{\reg}$ is a convex function. In Section 7 we proved that the hessian of $L_{\exp}$ is Lipschitz. In Section 8 we provide an approximate version of newton method for solving convex optimization problem which is more efficient under certain assumptions. In Section 9 we state our result of this paper and provide the algorithm for tackling the softmax regression problem.

Related Work

Since the explosion of LLM, there have been a lot of theoretical works about the computation of attention . Locality sensitive hashing (LSH) techniques have been employed in research to approximate attention. . Based on it, proposed KDEformer, an efficient approximation algorithm for the dot-product attention mechanism, with provable spectral norm bounds and superior performance on various pre-trained models. Recent research has investigated both static and dynamic approaches to attention computation . Additionally, delved into regularized hyperbolic regression problems involving exponential, cosh, and sinh functions. proposed randomized and deterministic algorithms to sparsify the attention matrix in large language models, achieving high accuracy with significantly reduced feature dimension.

Convergence and Optimization.

There have been works trying to understanding attention computation on optimization and convergence perspective . In practical attention models, adaptive methods often performs better than SGD. To understand this, showed that heavy-tailed distribution of the noise is one of the reason of the bad performance of SGD compared to adaptive methods, and provided new upper and lower bounds for convergence of adaptive methods under heavy-tailed noise in attention models. This answered the question of why adaptive methods performs better in attention models. explained why models sometimes attend to salient words and how the attention mechanism evolves throughout training, using a model property they defined, named Knowledge to Translate Individual Words (KTIW), which is learned early on from word co-occurrence statistics and later used to attend to input words while predicting the output. Recently, studied the regression problem inspired by the neural network with exponential activation function, and showed the convergence of a two-layer NN with large width (over-parameterized), while focused on solving regularized exp, cosh and sinh regression problems inspired by Attention computation. explored how transformers learn the co-occurrence structure of words by examining attention-based network size, depth, and complexity through experiments and mathematical analysis, showing that the embedding and self-attention layers encode topical structure with higher average inner product and pairwise attention between same-topic words.

Privacy and Security.

With the fast development of LLMs, the potential negative impact of abusing LLM has also been considered. To overcome this, without influencing the quality of the generated text, proposed a novel method to add watermark in LLM-generated text. The method needs no access to the parameters or API of the LLM. introduced a formal definition of near access-freeness (NAF) and develops generative model learning algorithms to ensure that the model outputs do not resemble copyrighted data by more than $k$-bits, with experiments on language (transformers) and image (diffusion) generative models demonstrating strong protection against sampling protected content.

2 Fast Linear Algebra

Let $t$ denote the target at one step of central (also mathematically called the complementarity gap). The $xs$ can viewed as reality. In the ideal case, they hope $xs=t$. However, this is unlikely to happen. They are using the potential function $\Phi(xs)=\sum_{i=1}^{n}\cosh(x_{i}s_{i}-t)$ to measure the difference between reality and target.

Sketching for Convex Optimization.

Sketching technique has been widely-used in optimization problems such as linear programming , empirical risk minimization , cutting plane method , computing John Ellipsoid , integral minimization problem , matrix completion , training over-parameterized neural tangent kernel regression , matrix sensing .

Technique Overview

Here in this section, we provide an overview of our techniques.

Recall the target function of our problem is in the form of

Calculating the Hessian of $L_{\exp}(x)$ directly is too complicated. To simplify this, we define two terms of $\alpha(x):=\langle\exp(Ax),\mathbf{1}_{n}\rangle$, $f(x):=\langle\exp(Ax),\mathbf{1}_{n}\rangle^{-1}\cdot\exp(Ax)$. Then to get the final Hessian to the loss functions, we calculate the Hessian step by step. To be specific, we divide the Hessian calculation into the following items:

Hessian of $\alpha(x)$ and $\alpha^{-1}(x)$;

After that, we notice a structured decomposition of Hessian of $L(x)$. We show that

where $B(x)$ is only function with $x$ and has no relation with respect to $i$ and $j$. In order to apply existing sparsification tool to boost the Hessian calculation (which is one of our main motivations), we construct specific decomposition to the two terms $B(x)$. We show that, $B$ can be viewed as sums of several rank-$1$ matrices and diagonal matrices.

Hessian is Positive Definite.

The key insight of this section lies in the analysis of volumetric barrier functions for solving semidefinite programming. . With the decomposition of the Hessian matrix for $L_{\exp}$, the next step is to bound it. To be specific, by dividing $B(x)$ in the way of low-rank parts and diagonal parts, we can lower and upper bound each segment of them. And by combining them, we can get the bound for $B(x)$,

This allows us to apply sparsification tool on $W$ to approximate the Hessian.

Lipschitz property for Hessian.

The key insight of this section lies in the analysis of previous analysis for recurrent neural networks . By the above calculation of Hessian, we divide the Hessian matrix to different segments. Now with the decomposition (to be specific, we divide the Hessian into low-rank parts and diagonal parts), we show Lipschitz property for each term. We first show Lipschitz property for the basic terms:

$\|\exp(Ax)\|_{2}\leq\sqrt{n}\cdot\exp(R^{2})$

$\|\exp(Ax)-\exp(Ay)\|_{2}\leq 2\sqrt{n}\cdot R\exp(R^{2})\cdot\|x-y\|_{2}$;

$|\alpha(x)-\alpha(y)|\leq\sqrt{n}\cdot\|\exp(Ax)-\exp(Ay)\|_{2}$;

$|\alpha(x)^{-1}-\alpha(y)^{-1}|\leq\beta^{-2}\cdot|\alpha(x)-\alpha(y)|$; (Later we will also prove an upper bound for $\beta^{-1}$, see Lemma 8.9)

$\|f(x)-f(y)\|_{2}\leq R_{f}\cdot\|x-y\|_{2}$. (Here $R_{f}$ is a function of $\beta^{-1},\exp(R^{2})$, see concrete definition in Lemma 7.2)

Then, following the decomposition of the Hessian matrix, we show the Lipschitz property for each of the divided terms (we use $G_{i}$ for $i\in 1,\dots,8$ to denote the terms) and combine them together to get the property of

for some small constant $\beta\in(0,0.1)$, which implies the Lipschitz property for the Hessian.

Approximated Newton Method with Sparsification Tool.

Newton method is a widely-used and traditional tool used in optimization questions. In many optimization applications, computing $\nabla^{2}L(x_{t})$ or $(\nabla^{2}L(x_{t}))^{-1}$ is quite expensive. Therefore, a natural motivation is to approximately formulate its Hessian or inverse of Hessian. In our setting, we want a faster implementation of Newton method. By above steps, we show our Hessian can be approximated by a matrix in the form of $A^{\top}DA$, where $D=W^{2}$ is a diagonal matrix. This inspires us to implement a standard tool that can generate a sparse matrix $\widetilde{D}$ such that

in near input-sparsity time of $A$. By this tool, we can reduce the time for Hessian calculation of each iteration to the time of $\widetilde{O}(\nnz(A)+d^{\omega})$. Here $\nnz(A)$ denotes the number of non-zero entries in matrix $A$. Let $\omega$ denote the exponent of matrix multiplication. Currently, $\omega\approx 2.373$ .

Preliminary

In this section, we provide the preliminaries used in our paper. In Section 4.1 we introduce the notations we use. In Section 4.2 we provide some basic facts for exact computation. In Section 4.3 we provide some tools for finding the bound of norms based on vectors. In Section 4.4 we provide some tools for finding the bound of norms related to matrices. In Section 4.5, we provide basic inequalities for psd matrices. In Section 4.6, we state several basic rules for calculus. In Section 4.7 we provide the regularization term $L_{\reg}$ and compute $\nabla L_{\reg}$ and $\nabla^{2}L_{\reg}$.

2 Basic Algebras

$\langle u,v\rangle=\langle u\circ v,{\bf 1}_{n}\rangle$

$\langle u\circ v,w\rangle=\langle u\circ v\circ w,{\bf 1}_{n}\rangle$

$u\circ v=v\circ u=\diag(u)\cdot v=\diag(v)\cdot u$

$u^{\top}(v\circ w)=v^{\top}(u\circ w)=w^{\top}(u\circ v)$

$u^{\top}\diag(v)w=v^{\top}\diag(u)w=u^{\top}\diag(w)v$

$\diag(u)\cdot\diag(v)\cdot{\bf 1}_{n}=\diag(u)v$

3 Basic Vector Norm Bounds

$\langle u,v\rangle\leq\|u\|_{2}\cdot\|v\|_{2}$ (Cauchy-Schwarz inequality)

$\|u\circ v\|_{2}\leq\|u\|_{\infty}\cdot\|v\|_{2}$

$\|u\|_{\infty}\leq\|u\|_{2}\leq\sqrt{n}\cdot\|u\|_{\infty}$

$\|u\|_{2}\leq\|u\|_{1}\leq\sqrt{n}\cdot\|u\|_{2}$

$\|\exp(u)\|_{\infty}\leq\exp(\|u\|_{\infty})\leq\exp(\|u\|_{2})$

Let $\alpha$ be a scalar, then $\|\alpha\cdot u\|_{2}=|\alpha|\cdot\|u\|_{2}$

For any $\|u-v\|_{\infty}\leq 0.01$, we have $\|\exp(u)-\exp(v)\|_{2}\leq\|\exp(u)\|_{2}\cdot 2\|u-v\|_{\infty}$

For all the other facts we omit the details. We will only prove the last fact.

where the 1st step follows from definition of $\circ$ operation and $\exp()$, the 2nd step follows from $\|u\circ v\|_{2}\leq\|u\|_{\infty}\cdot\|v\|_{2}$, the 3rd step follows from $|\exp(x)-1|\leq 2x$ for all $x\in(0,0.1)$.

4 Basic Matrix Norm Bounds

If $U\preceq\alpha\cdot V$, then $\|U\|\leq\alpha\cdot\|V\|$

For any vector $v$, we have $\|Uv\|_{2}\leq\|U\|\cdot\|v\|_{2}$.

5 Basic PSD

$uu^{\top}\preceq\|u\|_{2}^{2}\cdot I_{n}$.

$\diag(u\circ u)\preceq\|u\|_{2}^{2}\cdot I_{n}$

$uv^{\top}+vu^{\top}\preceq uu^{\top}+vv^{\top}$

$uv^{\top}+vu^{\top}\succeq-(uu^{\top}+vv^{\top})$

$(v\circ u)(v\circ u)^{\top}\preceq\|v\|_{\infty}^{2}uu^{\top}$

6 Basic Derivative Rules

Let $f$ denote a differentiable function.

7 Regularization

Softmax Regression Loss

In this section, we provide detailed computation for $\nabla L_{\exp}$ and $\nabla^{2}L_{\exp}$. In Section 5.1, we define $f(x)$ and $\alpha(x)$ to simplify the computation for $\nabla L_{\exp}$ and $\nabla^{2}L_{\exp}$. In Section 5.2, we compute $\nabla L_{\exp}$ step by step. In Section 5.3, we define the gradient of Loss function and also prove the Lipschitz property for gradient. In Section 5.4-5.8, we compute $\nabla^{2}L_{\exp}$ step by step. To be specific, in Section 5.4, we compute $\nabla^{2}\exp(Ax)$; in Section 5.5, we compute $\nabla^{2}\alpha(x)$; in Section 5.6, we compute $\nabla^{2}\alpha(x)^{-1}$; in Section 5.7, we compute $\nabla^{2}f(x)$; in Section 5.8, we compute $\nabla^{2}L_{\exp}$. In Section 5.9, we provide some result to aid the computation in Section 5.10. In Section 5.10, we split $\nabla^{2}L_{\exp}$ into several low rank matrices and diagonal matrices.

We define function softmax $f$ as follows

For any vector $b$, $0\preceq(b\circ f(x))(b\circ f(x))^{\top}\preceq\|b\|_{\infty}^{2}f(x)f(x)^{\top}\preceq\|b\|_{\infty}^{2}I_{n}$

For any vector $b$, $\diag(b\circ b)\preceq\|b\|_{\infty}^{2}I_{n}$

$0\preceq\diag(f(x))\preceq\|f(x)\|_{\infty}I_{n}\preceq\|f(x)\|_{2}I_{n}$.

$0\preceq\diag(f(x)\circ f(x))\preceq\|f(x)\|_{\infty}^{2}I_{n}\preceq\|f(x)\|_{2}I_{n}$.

The proofs are very straightforward, so we omitted the details here. ∎

For convenient, we define two helpful notations $\alpha$ and $c$

Then, we can rewrite $f(x)$ (see Definition 5.1) and $L_{\exp}(x)$ (see Definition 5.3) as follows

$L_{\exp}(x)=0.5\cdot\|\alpha(x)^{-1}\cdot\exp(Ax)-b\|_{2}^{2}$.

$L_{\exp}(x)=0.5\cdot\|f(x)-b\|_{2}^{2}$.

Then we can rewrite $L_{\exp}(x)$ (see Definition 5.3) as follows

2 Gradient

Let $\alpha(x)$ be defined in Definition 5.4.

Let $L_{\exp}(x)$ be defined in Definition 5.3.

Proof of Part 1. For each $i\in[d]$, we have

where the 1st step follows from simple algebra, the 2nd step follows from Fact 4.6, the 3rd step follows from simple algebra.

Proof of Part 2. It trivially follows from arguments in Part 1.

where the 1st step follows from Definition of $f$, 2nd step follows from differential chain rule, the 3rd step follows from the result from Part 2 and Part 3, the forth step follows from definition of $f$ (see Definition 5.1).

where the 1st step follows from extracting $A_{*,i}$, the 2nd step follows from result of Part 4, the 3rd step follows from simple algebra, the last step follows from simple algebra.

where the 1st step follows from extracting $A_{*,i}$, the 2nd step follows from result of Part 4, the 3rd step follows from simple algebra, the 4th step follows from simple algebra, the last step follows from simple algebra.

3 Definition of Gradient

In this section, we use $g(x)$ to denote the gradient of $L_{\exp}(x)$.

Let $L_{\exp}(x)$ be defined as Definition 5.3.

Equivalently, for each $i\in[d]$, we define

$\|f(x)-f(y)\|_{2}\leq R_{f}\cdot\|x-y\|_{2}$

$\|c(x)-c(y)\|_{2}\leq R_{f}\cdot\|x-y\|_{2}$

Let $R_{\infty}\in(0,2]$ be parameter such that

where the 1st step follows from the definition of $g_{1}$, the 2nd step follows from adding some terms $-f(y)\langle c(x),f(x)\rangle+f(y)\langle c(x),f(x)\rangle-f(y)\langle c(y),f(x)\rangle+f(y)\langle c(y),f(x)\rangle$ and $\|a+b\|_{2}\leq\|a\|_{2}+\|b\|_{2}$(Fact 4.3).

where the 1st step follows from $\|\alpha a\|_{2}\leq|\alpha|\|a\|_{2}$(Fact 4.3), the 2nd step follows from $\langle a,b\rangle\leq\|a\|_{2}\|b\|_{2}$(Fact 4.3), the 3rd step follows from the definition of $R_{f}$.

where the 1st step follows from $\|\alpha a\|_{2}\leq|\alpha|\|a\|_{2}$(Fact 4.3), the 2nd step follows from $\langle a,b\rangle\leq\|a\|_{2}\|b\|_{2}$(Fact 4.3), the 3rd step follows from the definition of $R_{f}$.

the 2nd step follows from $\langle a,b\rangle\leq\|a\|_{2}\|b\|_{2}$(Fact 4.3), the 3rd step follows from the definition of $R_{f}$.

Combining three terms together, we complete the proof.

this step follows from adding terms $-\diag(f(x))c(y)+\diag(f(x))c(y)$.

Combining two terms together, then we complete the proof. ∎

It follows from combining Part 1 and Part 2.

4 Hessian Calculations: Step 1, Hessian of exp⁡(A​x)\exp(Ax)

where the 1st step is an expansion of the Hessian, the 2nd step follows from the differential chain rule, the 3rd step extracts the matrix $A_{*,i}$ with constant entries out of the derivative, and the last step also follows from the chain rule.

where the 1st step is an expansion of the Hessian, the 2nd and 3rd steps follow from the differential chain rule, the 3rd step follows from simple algebra.

5 Hessian Calculations: Step 2, Hessian of α⁡(x)\alpha(x)

Let $\alpha(x)$ be defined as Definition 5.4.

where the 1st step follows from the expansion of hessian, the 2nd step follows from Part 3 of Lemma 5.6, the 3rd step follows from simple algebra, and the last step follows from Fact 4.1.

where the 1st step follows from the expansion of hessian, the 2nd step follows from Part 2 of Lemma 5.6, the 3rd step follows from simple algebra, the last step follows from Fact 4.1.

6 Hessian Calculations: Step 3, Hessian of α​(x)−1\alpha(x)^{-1}

Let $\alpha(x)$ be defined as Definition 5.4

where the 1st step follows from the expansion of hessian, the 2nd step follows from Part 3 of Lemma 5.6, the 3rd step follows from differential chain rule, the 4th step follows from simple algebra, the last step follows from Fact 4.1.

where the 1st step follows from the expansion of hessian, the 2nd step follows from Part 3 of Lemma 5.6, the 3rd step follows from differential chain rule, the 4th step follows from basic differential rule, the 5th step step follows from simple algebra, the last step follows from Fact 4.1. ∎

7 Hessian Calculations: Step 4, Hessian of f⁡(x)f(x)

Let $f(x)=\langle\exp(Ax),{\bf 1}_{n}\rangle^{-1}\exp(Ax)$ (see Definition 5.1).

where the 1st step follows from the expansion of hessian, the 2nd step follows from Part 4 of Lemma 5.6 and differential chain rule.

where the 1st step follows from the expansion of hessian, the 2nd step follows from Part 4 of Lemma 5.6, the 3rd step follows from differential chain rule.

8 Hessian Calculations: Step 5, Hessian of Lexp​(x)L_{\exp}(x)

where the 1st step follows from the expansion of hessian, the 2nd step follows from differential chain rule.

where the 1st step follows from the expansion of hessian, the 2nd step follows from differential chain rule, the 3rd step is a simplification of step 2 by applying notations $\alpha$ (Definition 5.4) and $c$ ( Definition 5.5 ).

9 Helpful Lemma

The goal of this section to prove Lemma 5.14. We remark that in this lemma, we can replace $f(x)$ by any vector. However, for easy of presentation, we use $f(x)$.

where the 1st step follows from Fact 4.2, the 2nd step follows from Fact 4.2.

where the 1st step follows from $a^{\top}b=\langle a,b\rangle$ (Fact 4.1), the 2nd step follows from $\langle a,b\rangle=a^{\top}b$ (Fact 4.1).

where the 1st step follows from $\langle a,b\rangle=a^{\top}b$ (Fact 4.1), the 2nd step follows from Fact 4.2, the 3rd step follows from $a^{\top}b=\langle a,b\rangle$ (Fact 4.1), the last step follows from Fact 4.2.

where the 1st step follows from $\langle a,b\rangle=a^{\top}b$ (Fact 4.1), the 2nd step follows from Fact 4.2, the 3rd step follows from $f(x)^{\top}A_{*,j}=\langle f(x),A_{*,j}\rangle$ (Fact 4.1) is a scalar.

where the 1st step follows from $\langle a,b\rangle=a^{\top}b$ (Fact 4.1), the 2nd step follows from Fact 4.2, the 3rd step follows from $a^{\top}b=b^{\top}a$ (Fact 4.1), the last step follows from $a^{\top}b=b^{\top}a$ (Fact 4.1).

where the 1st step follows from Fact 4.2, the 2nd step follows from $\langle a,b\rangle=a^{\top}b$ (Fact 4.1), the 3rd step follows from $f(x)^{\top}A_{*,j}=\langle f(x),A_{*,j}\rangle$ is a scalar (Fact 4.1).

where the 1st step follows from $a^{\top}b=\langle a,b\rangle$ (Fact 4.1), the 2nd step follows from Fact 4.1, the 3rd step follows from $\langle a,b\rangle=a^{\top}b$ (Fact 4.1), the 4th step follows from Fact 4.2, the last step follows from Fact 4.2.

where the 1st step follows from $a^{\top}b=b^{\top}a$ (Fact 4.1), the 2nd step follows from $\langle a,b\rangle=a^{\top}b$ (Fact 4.1), the 3rd step follows from $a^{\top}b=b^{\top}a$ (Fact 4.1), the 4th step follows from $A_{*,i}^{\top}f(x)=\langle A_{*,i},f(x)\rangle$ is a scalar (Fact 4.1), the 5th step step follows from $f(x)^{\top}A_{*,j}=\langle f(x),A_{*,j}\rangle$ is a scalar (Fact 4.1), the last step follows from $a^{\top}b=\langle a,b\rangle$ (Fact 4.1).

where the 1st step follows from $a^{\top}b=\langle a,b\rangle$ (Fact 4.1), the 2nd step follows from Fact 4.1 , the 3rd step follows from $\langle a,b\rangle=a^{\top}b$ (Fact 4.1), the 4th step follows from Fact 4.2, the last step follows from Fact 4.2. ∎

10 Decomposing B1​(x)B_{1}(x), B2​(x)B_{2}(x) and B⁡(x)B(x) into Low Rank Plus Diagonal

where the 1st step follows from the definition of $B_{1}(x)$, the 2nd step follows from $(A+B)^{\top}=A^{\top}+B^{\top}$, the 3rd step follows from simple algebra, the last step follows from Lemma 5.14.

Thus, by extracting $A_{*,i}^{\top}$ and $A_{*,j}$, we have:

where the 1st step follows from definition of $B_{2}(x)$, the 2nd step follows from simple algebra, the 3rd step follows from simple algebra, the last step follows from Lemma 5.14.

By extracting $A_{*,i}^{\top}$ and $A_{*,j}$, we have

By combining all the above equations, we have

Hessian is Positive Definite

In this section, we prove that $\nabla^{2}L\succeq 0$ and thus $L$ is convex. In Section 6.1, we find the lower bound of $B(x)$. To be specific, we split $B(x)$ into several terms and find their lower bounds separately. In Section 6.2, we use the result of Section 6.1 to prove that lower bound of $\nabla^{2}L\succeq 0$ and thus $L$ is convex.

Let $B_{\rank}^{1}$, $B_{\rank}^{2}$ be defined as Definition 6.1.

Let $B_{\diag}^{1}$, $B_{\diag}^{2}$ be defined as Definition 6.1.

Part 5. If $\|b\|_{1}\leq 1$ and $\|f(x)\|_{1}\leq 1$, then we have

Recall that in Definition 6.1, we split $B(x)$ into four terms

where $B_{\rank}^{i}$ and $B_{\diag}^{i}$ are defined as

On one hand, we can lower bound the coefficient, we have

where the 1st step follows from Fact 4.5, , the last step follows from Fact 4.5.

where the 1st step follows from Fact 4.5 , the 2nd step follows from Fact 4.5 .

where the 1st step follows from simple algebra, the 2nd step follows from simple algebra, the last step follows from Fact 4.5.

Proof of $B(x)$. It trivially follows from

2 Lower bound on Hessian

The goal of this section is to prove Lemma 6.3.

Let $L_{\exp}(x)$ be defined as Definition 5.3.

Let $L_{\reg}(x)$ be defined as Definition 4.8.

Let $\sigma_{\min}(A)$ denote the minimum singular value of $A$.

Part 1. If all $i\in[n]$, $w_{i}^{2}\geq 4+l/\sigma_{\min}(A)^{2}$, then

Part 2. If all $i\in[n]$, $w_{i}^{2}\geq 100+l/\sigma_{\min}(A)^{2}$, then

By applying Lemma 5.13 and Lemma 5.15, we have

Thus, by applying Lemma 4.9, Eq. (4) can be written as

where the 3rd step follows from $w_{\min}^{2}\geq 4+l/\sigma_{\min}(A)^{2}$, the last step follows from simple algebra.

Since $D$ is positive definite, then we have

Thus, Hessian is positive definite forever and thus the function is convex. ∎

Hessian is Lipschitz

In this section, we find the upper bound of $\|\nabla^{2}L(x)-\nabla^{2}L(y)\|$ and thus proved that $\nabla^{2}L$ is lipschitz. In Section 7.2, we prove that some basic terms satisfy the property of Lipschitz. In Section 7.3, we provide a sketch of how we find the bound of $\|\nabla^{2}L(x)-\nabla^{2}L(y)\|$, to be specific, we split $\|\nabla^{2}L(x)-\nabla^{2}L(y)\|$ into 8 terms and state that all these terms can be bound by using $\|f(x)-f(y)\|$. In Section 7.4, we use $\|f(x)-f(y)\|$ to bound the first term. In Section 7.5, we use $\|f(x)-f(y)\|$ to bound the second term. In Section 7.6, we use $\|f(x)-f(y)\|$ to bound the third term. In Section 7.7, we use $\|f(x)-f(y)\|$ to bound the fourth term. In Section 7.8, we use $\|f(x)-f(y)\|$ to bound the fifth term. In Section 7.9, we use $\|f(x)-f(y)\|$ to bound the sixth term. In Section 7.10, we use $\|f(x)-f(y)\|$ to bound the seventh term. In Section 7.11, we use $\|f(x)-f(y)\|$ to bound the last term.

$\langle\exp(Ax),{\bf 1}_{n}\rangle\geq\beta$ and $\langle\exp(Ay),{\bf 1}_{n}\rangle\geq\beta$

where the 1st step follows definition of $G_{i}$ and matrix spectral norm, the 2nd step follows from $\|A\|\leq R$, the 3rd step follows from Lemma 7.3, the 4th step follows from Lemma 7.2, and the last step follows from simple algebra. ∎

2 A Core Tool: Lipschitz Property for Several Basic Functions

$\langle\exp(Ax),{\bf 1}_{n}\rangle\geq\beta$

$\langle\exp(Ay),{\bf 1}_{n}\rangle\geq\beta$

Let $R_{f}:=\beta^{-2}n^{1.5}\exp(3R^{2})$

Let $\alpha(x)$ be defined as Definition 5.4

Part 0. $\|\exp(Ax)\|_{2}\leq\sqrt{n}\exp(R^{2})$

Part 1. $\|\exp(Ax)-\exp(Ay)\|_{2}\leq 2\sqrt{n}R\exp(R^{2})\cdot\|x-y\|_{2}$

Part 2. $|\alpha(x)-\alpha(y)|\leq\sqrt{n}\cdot\|\exp(Ax)-\exp(Ay)\|_{2}$

Part 3. $|\alpha(x)^{-1}-\alpha(y)^{-1}|\leq\beta^{-2}\cdot|\alpha(x)-\alpha(y)|$

Part 4. $\|f(x)-f(y)\|_{2}\leq R_{f}\cdot\|x-y\|_{2}$

Part 5. $\|c(x)-c(y)\|_{2}\leq R_{f}\cdot\|x-y\|_{2}$

Part 6. $\|g(x)-g(y)\|_{2}\leq 16\cdot R\cdot R_{f}\cdot\|x-y\|_{2}$

where the first step follows from Fact 4.3, the second step follows from Fact 4.3, the third step follows from Fact 4.3, and the last step follows from $\|A\|\leq R$ and $\|x\|_{2}\leq R$.

where the 1st step follows from $\|A(y-x)\|_{\infty}<0.01$ and Fact 4.3, the 2nd step follows from Part 0, the 3rd step follows from Fact 4.3, the 4th step follows from Fact 4.4, the last step follows from $\|A\|\leq R$.

where the 1st step follows from the definition of $\alpha(x)$, the 2nd step follows from Cauchy-Schwarz inequality (Fact 4.3).

where the 1st step follows from simple algebra, the 2nd step follows from $\alpha(x),\alpha(y)\geq\beta$.

where the 1st step follows from the definition of $f(x)$ and $\alpha(x)$, the 2nd step follows from triangle inequality, the 3rd step follows from $\|\alpha A\|\leq|\alpha|\|A\|$(Fact 4.4).

where the 1st step follows from $\alpha(x)\geq\beta$, the 2nd step follows from Part 1.

For the second term in the above, we have

where the 1st step follows from the result of Part 3, the 2nd step follows from Part 0, the 3rd step follows from the result of Part 2, the 4th step follows from Part 1, and the last step follows from simple algebra.

Combining Eq. (7.2) and Eq. (7.2) together, we have

where the 1st step follows from the bound of the first term and the second term, the 2nd step follows from $\beta^{-1}\geq 1$ and $n>1$ trivially, the 3rd step follows from simple algebra.

the first step follows from the definition of $c(x)$, the last step follows from Part 4 and definition of $R_{f}$. Proof of Part 6.

where the second step follows from $\|A\|\leq R$.

3 Summary of Eight Steps

$G_{1}=\|f(x)\|_{2}^{2}f(x)f(x)^{\top}-\|f(y)\|_{2}^{2}f(y)f(y)^{\top}$

$G_{2}=\langle f(x),b\rangle f(x)f(x)^{\top}-\langle f(y),b\rangle f(y)f(y)^{\top}$

$G_{3}=\langle f(x),f(x)\rangle\diag(f(x))-\langle f(y),f(y)\rangle\diag(f(y))$

$G_{4}=\langle f(x),b\rangle\diag(f(x))-\langle f(y),b\rangle\diag(f(y))$

$G_{5}=\diag(f(x)\circ(f(x)-b))-\diag(f(y)\circ(f(y)-b))$

$G_{6}=\diag(f(x)\circ f(x))-\diag(f(y)\circ f(y))$

$G_{7}=f(x)(f(x)\circ b)^{\top}-f(y)(f(y)\circ b)^{\top}$

$G_{8}=(f(x)\circ b)f(x)^{\top}-(f(y)\circ b)f(y)^{\top}$

The proof directly follows from applying Lemma 7.4, Lemma 7.5, Lemma 7.6, Lemma 7.7, Lemma 7.8, Lemma 7.9, Lemma 7.10, Lemma 7.11. ∎

4 Lipschitz Calculations: Step 1. Lipschitz for Matrix Function ‖f⁡(x)‖22​f​(x)​f​(x)⊤\|f(x)\|_{2}^{2}f(x)f(x)^{\top}

$G_{1}=\|f(x)\|_{2}^{2}f(x)f(x)^{\top}-\|f(y)\|_{2}^{2}f(y)f(y)^{\top}$

Let us only prove for $G_{1,1}$, the others are similar,

where the 1st step follows from Fact 4.4, the 2nd step follows from $|\langle a,b\rangle|\leq\|a\|_{2}\|b\|_{2}$ (Fact 4.3), the 3rd step follows from $aa^{\top}\preceq\|a\|_{2}^{2}I_{n}$(Fact 4.3), the last step follows from $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$ (Lemma 5.2).

It is obvious that for each $i\in$, we have

where the last step follows from $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$. ∎

$G_{2}:=\langle f(x),b\rangle f(x)f(x)^{\top}-\langle f(y),b\rangle f(y)f(y)^{\top}$

Since $G_{2,1},G_{2,2},G_{2,3}$ are similar, we only have to bound $\|G_{2,1}\|$:

where the 1st step follows from the definition of $G_{2,1}$, the 2nd step follows from simple algebra, the 3rd step follows from Fact 4.4, the 4th step follows from $\|ab^{\top}\|\leq\|a\|_{2}\|b\|_{2}$(Fact 4.4), the 5th step follows from $\langle a,b\rangle\leq\|a\|_{2}\|b\|_{2}$(Fact 4.3), the last step follows from $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$.

6 Lipschitz Calculations: Step 3. Lipschitz for Matrix Function f⁡(x)​f​(x)⊤​diag⁡(f⁡(x))f(x)f(x)^{\top}\diag(f(x))

$G_{3}:=\langle f(x),f(x)\rangle\diag(f(x))-\langle f(y),f(y)\rangle\diag(f(y))$

Since $G_{3,1},G_{3,2},G_{3,3}$ are similar, we only need to bound $\|G_{3,1}\|$:

where the 1st step follows from the definition of $G_{3,1}$, the 2nd step follows from simple algebra, the 3rd step follows from $\|\alpha A\|\leq|\alpha|\|A\|$(Fact 4.4), $\langle a,b\rangle\leq\|a\|_{2}\|b\|_{2}$(Fact 4.3), and $\|ab\|\leq\|a\|\|b\|$(Fact 4.4), the 4th step follows from $\|\diag(f(x))\|=\|f(x)\|_{2}$, the last step follows from $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$ (Fact 4.3).

where the 1st step follows from the definition of $G_{3}$, the 2nd step follows from Fact 4.4, the last step follows from the bound of $\|G_{3,1}\|$,$\|G_{3,2}\|$ and $\|G_{3,3}\|$. ∎

7 Lipschitz Calculations: Step 4. Lipschitz for Matrix Function ⟨f⁡(x),b⟩​diag⁡(f⁡(x))\langle f(x),b\rangle\diag(f(x))

$G_{4}:=\langle f(x),b\rangle\diag(f(x))-\langle f(y),b\rangle\diag(f(y))$

Since $G_{4,1}$ and $G_{4,2}$ are similar, we only need to bound $\|G_{4,1}\|$:

where the 1st step follows from the definition of $G_{4,1}$, the 2nd step follows from simple algebra, the 3rd step follows from $\|ab\|\leq\|a\|\|b\|$(Fact 4.4) and , the 4th step follows from $\|\diag(x)\|\leq\|x\|_{\infty}\leq\|x\|_{2}$(Fact 4.3), the last step follows from $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$(Fact 4.3).

where the 1st step follows from the definition of $G_{4}$, the 2nd step follows from Fact 4.4, the last step follows from the bound of $\|G_{4,1}\|$ and $\|G_{4,2}\|$. ∎

8 Lipschitz Calculations: Step 5. Lipschitz for Matrix Function diag⁡(f⁡(x)∘(f⁡(x)−b))\diag(f(x)\circ(f(x)-b))

$G_{5}:=\diag(f(x)\circ(f(x)-b))-\diag(f(y)\circ(f(y)-b))$

where the 1st step follows from the definition of $G_{5,1}$, the 2nd step follows from Fact 4.2, the 3rd step follows from $\|ab\|\leq\|a\|\|b\|$(Fact 4.4), the 4th step follows from $\|\diag(a)\|\leq\|a\|_{\infty}\leq\|a\|_{2}$(Fact 4.3), the last step follows from $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$.

where the 1st step follows from the definition of $G_{5,2}$, the 2nd step follows from Fact 4.2, the 3rd step follows from Fact 4.4, the 4th step follows from $\|\diag(a)\|\leq\|a\|_{\infty}\leq\|a\|_{2}$(Fact 4.3), the last step follows from $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$.

where the 1st step follows from the definition of $G_{5}$, the 2nd step follows from Fact 4.2, the 3rd step follows from the bound of $\|G_{5,1}\|$ and $\|G_{5,2}\|$. ∎

9 Lipschitz Calculations: Step 6. Lipschitz for Matrix Function diag⁡(f⁡(x)∘f⁡(x))\diag(f(x)\circ f(x))

$G_{6}:=\diag(f(x)\circ f(x))-\diag(f(y)\circ f(y))$

Since, $G_{6,1}$ and $G_{6,2}$ are similar, we only need to bound $\|G_{6,1}\|$:

where the 1st step follows from the definition of $G_{6,1}$, the 2nd step follows from Fact 4.2, the 3rd step follows from $\|ab\|\leq\|a\|\|b\|$(Fact 4.4) and $\|\diag(a)\|\leq\|a\|_{\infty}\leq\|a\|_{2}$(Fact 4.3), the last step follows from $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$.

where the 1st step follows from the definition of $G_{6}$, the 2nd step follows from Fact 4.4, the last step follows from the bound of $\|G_{6,1}\|$ and $\|G_{6,2}\|$. ∎

10 Lipschitz Calculations: Step 7. Lipschitz for Matrix Function f⁡(x)​(f⁡(x)∘b)⊤f(x)(f(x)\circ b)^{\top}

$G_{7}:=f(x)(f(x)\circ b)^{\top}-f(y)(f(y)\circ b)^{\top}$

Since $G_{7,1}$ and $G_{7,2}$ are similar, we only need to bound $\|G_{7,1}\|$:

where the 1st step follows from the definition of $G_{7,1}$, the 2nd step follows from simple algebra, the 3rd step follows from $\|ab^{\top}\|\leq\|a\|_{2}\|b\|_{2}$(Fact 4.4) and $\|a^{\top}\|_{2}=\|a\|_{2}$, the last step follows from $\|a\circ b\|_{2}\leq\|a\|_{\infty}\|b\|\leq\|a\|_{2}\|b\|_{2}$(Fact 4.3) and $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$.

where the 1st step follows from the definition of $G_{7}$, the 2nd step follows from Fact 4.4, the last step follows from the bound of $\|G_{7,1}\|$ and $\|G_{7,2}\|$. ∎

11 Lipschitz Calculations: Step 8. Lipschitz for Matrix Function (f⁡(x)∘b)​f​(x)⊤(f(x)\circ b)f(x)^{\top}

$G_{8}:=(f(x)\circ b)f(x)^{\top}-(f(y)\circ b)f(y)^{\top}$

Since $G_{8,1}$ and $G_{8,2}$ are similar, we only need to bound $\|G_{8,1}\|$:

where the 1st step follows from the definition of $G_{8,1}$, the 2nd step follows from simple algebra, the 3rd step follows from $\|ab^{\top}\|\leq\|a\|_{2}\|b\|_{2}$(Fact 4.4), the 4th step follows from $\|a\circ b\|_{2}\leq\|a\|_{\infty}\|b\|\leq\|a\|_{2}\|b\|_{2}$(Fact 4.3), the last step follows from $\|f(x)\|_{2}\leq\|f(x)\|_{1}\leq 1$ (Lemma 5.2).

where the 1st step follows from the definition of $G_{8}$, the 2nd step follows from Fact 4.4, the last step follows from the bound of $\|G_{8,1}\|$ and $\|G_{8,2}\|$. ∎

Approximate Newton Method

In this section, we provide an approximate version of the newton method for convex optimization. In Section 8.1, we state some assumptions of the traditional newton method and the exact update rule of the traditional algorithm. In Section 8.2, we provide the approximate update rule of the approximate newton method, we also implement a tool for compute the approximation of $\nabla^{2}L$ and use some lemmas from to analyze the approximate newton method. In Section 8.3, we prove a lower bound on $\beta$. In Section 8.4, we prove an upper bound on $M$.

Here in this section, we focus on the local convergence of the Newton method. We consider the following target function

$\nabla^{2}L(x^{*})\succeq l\cdot I_{d}$.

Hessian is $M$-Lipschitz. If there exists a positive scalar $M>0$ such that

Good Initialization Point. Let $x_{0}$ denote the initialization point. If $r_{0}:=\|x_{0}-x_{*}\|_{2}$ satisfies

We define gradient and Hessian as follows

2 Approximate of Hessian and Update Rule

In many real-world tasks, it is very hard and expensive to compute exact $\nabla^{2}L(x_{t})$ or $(\nabla^{2}L(x_{t}))^{-1}$. Thus, it is natural to consider the approximated computation of the gradient and Hessian. The computation is defined as

In order to get the approximated Hessian $\widetilde{H}(x_{t})$ efficiently, here we state a standard tool (see Lemma 4.5 in ).

Note that, $\omega$ denotes the exponent of matrix multiplication, currently $\omega\approx 2.373$ .

Following the standard of Approximate Newton Hessian literature , we consider the following.

Loss Function $L$ is $(l,M)$-good (see Definition 8.1).

Let $\epsilon_{0}\in(0,0.1)$ (see Definition 8.4).

Let $T$ denote the total number of iterations of the algorithm, to apply Lemma 8.7, we will need the following induction hypothesis lemma. This is very standard in the literature, see .

For each $i\in[t]$, we define $r_{i}:=\|x_{i}-x^{*}\|_{2}$. If the following condition hold

$\epsilon_{0}=0.01$ (see Definition 8.4 for $\epsilon_{0}$)

$r_{i}\leq 0.4\cdot r_{i-1}$, for all $i\in[t]$

$M\cdot r_{i}\leq 0.1l$, for all $i\in[t]$ (see Definition 8.1 for $M$)

3 Lower bound on β\beta

Let $\beta$ be lower bound on $\langle\exp(Ax),{\bf 1}_{n}\rangle$

4 Upper bound on MM

Let $H$ denote the hessian of loss function $L$.

$\|H(x)-H(y)\|\leq\beta^{-2}n^{1.5}\exp(20R^{2})\cdot\|x-y\|_{2}$ (Lemma 7.1)

Main Result

Define $f(x):=\langle\exp(Ax),\mathbf{1}_{n}\rangle^{-1}\exp(Ax)$.

Define $x^{*}$ as the optimal solution of

It holds that $b\geq{\bf 0}_{n}$, and $\|b\|_{1}\leq 1$.

It holds that $w_{i}^{2}\geq 100+l/\sigma_{\min}(A)^{2}$ for all $i\in[n]$

Let $x_{0}$ denote an initial point for which it holds that $M\|x_{0}-x^{*}\|_{2}\leq 0.1l$.

Here $\omega$ denote the exponent of matrix multiplication. Currently $\omega\approx 2.373$ .

It follows from combining Lemma 6.3, Lemma 8.8, Lemma 8.5, Lemma 7.1 and Lemma 8.7.

Proof of Number of Iterations. After $T$ iterations, we have

By choice of $T$, we get the desired bound. The failure probability is following from union bound over $T$ iterations. ∎

References