SlowGuess/NeurIPS_2023_OCR huggingface dataset:
for all  $S_{i}\subset [n]$  with  $|S_i|\leq k$
- for all  $X \subset [n]$ ,  $X \subseteq S_i \cup \{i\}$  and  $|X \cap S_i| \leq |S_i|$ ,
$$
\left| \left(\widetilde {f} _ {S _ {i}, i} (X)\right) - \widetilde {f} _ {S _ {i}, i} (S _ {i}) - \left(f _ {S _ {i}, i} (X) - f _ {S _ {i}, i} (S _ {i})\right) \right| \leq \epsilon / (2 n)
$$
Let  $\mathsf{opt}_i$  denote the optimal cost that can be achieved by using  $f$  in  $i$ -th iteration. If the greedy algorithm use  $f$  to find the solution has performance at least  $(1 - 1 / e)\cdot \mathsf{opt}_i$ , then using  $\widetilde{f}$ , we can obtain a solution that has performance at least  $(1 - 1 / e)\cdot \mathsf{opt}_i - \epsilon$
Proof. The proof follows from Lemma D.18.  $\square$
Next, we explain how to add items into sets based on exact values.
Definition D.20 (Expanding items based on exact value). If the following conditions hold
- Let  $S_{i} \subseteq [i - 1]$ .
- Let  $f_{S_i,i} : 2^{[i]} \to \mathbb{R}$ .
Then, we can define  $S_{i+1}$  as follows
- If  $|S_i| = k$  then  $S_{i+1} = S_i \cup \{i\} \backslash u$  where  $u = \arg \max_{v \in S_i \cup \{i\}} f_{S_i,i}(S_i \cup \{i\} \backslash v) - f_{S_i,i}(S_i)$ .
- If  $|S_i| < k$ , then  $S_{i+1} = S_i \cup \{i\}$ .
Remark D.21. We remark that  $u = \arg \max_{v \in S_i \cup \{i\}} f_{S_i,i}(S_i \cup \{i\} \backslash v) - f_{S_i,i}(S_i)$  is the same as  $u = \arg \max_{v \in S_i \cup \{i\}} f_{S_i,i}(S_i \cup \{i\} \backslash v)$ . For the convenience of discussion and analysis, we will switch to using both cases in different places.
Here, we explain how to add items into sets via approximate values.
Definition D.22 (Expanding items based on approximate value). If the following conditions hold
- Let  $S_{i} \subseteq [i - 1]$ .
- Let  $\widetilde{f}_{S_i,i}:2^{[i]}\to \mathbb{R}$
- If  $|S_i| = k$  then  $S_{i+1} = S_i \cup \{i\} \backslash u$  where  $u = \arg \max_{v \in S_i \cup \{i\}} \widetilde{f}_{S_i,i}(S_i \cup \{i\} \backslash v) - \widetilde{f}_{S_i,i}(S_i)$ .
- If  $|S_i| < k$ , then  $S_{i+1} = S_i \cup \{i\}$ .
D.11 Universal Conditions
We state several definitions for both the exact function and approximation function.
Definition D.23 (Universal Monotone Condition (for exact function)). We say  $f$  has universal monotone condition, if for all  $i \in [n]$  for all  $S_{i} \subseteq [i - 1]$ , we have
$$
f _ {S _ {i}, i} (X) \geq f _ {S _ {i}, i} (Y), \quad \forall Y \subset X
$$
Definition D.24 (Universal Monotone Condition (for approximate function)). We say  $\widetilde{f}$  has universal monotone condition, if for all  $i\in [n]$  for all  $S_{i}\subseteq [i - 1]$ , we have
$$
\widetilde {f} _ {S _ {i}, i} (X) \geq \widetilde {f} _ {S _ {i}, i} (Y), \forall Y \subset X
$$
Definition D.25 (Universal Dynamic Condition 1(for exact function)). We say  $f$  has universal dynamic condition 1(for exact function), if for all  $i \in [n]$  for all  $S_i \subseteq [i-1]$ ,  $S_{i-1} \subseteq [i-2]$ ,  $|S_i \setminus S_{i-1}| \leq 1$ , we have
$$
f _ {S _ {i}, i} \left(S _ {i}\right) \geq (1 - \theta) \cdot f _ {S _ {i - 1}, i - 1} \left(S _ {i}\right)
$$
Definition D.26 (Universal Dynamic Condition 1 (for approximate function)). We say  $\widetilde{f}$  has universal dynamic condition 1 (for approximate function), if for all  $i\in [n]$  for all  $S_{i}\subseteq [i - 1]$ ,  $S_{i - 1}\subseteq [i - 2]$ ,  $|S_{i}\backslash S_{i - 1}|\leq 1$ , we have
$$
\widetilde {f} _ {S _ {i}, i} \left(S _ {i}\right) \geq (1 - \theta) \cdot \widetilde {f} _ {S _ {i - 1}, i - 1} \left(S _ {i}\right)
$$
We define  $\mathrm{opt}_i$  as follows:
Definition D.27. Let  $k$  denote the budget length. For each  $i \in [n]$ , we define  $\mathrm{opt}_i$  as
$$
\max  \left\{f _ {X, i} (Y) \mid X \subseteq [ i - 1 ], Y \subseteq [ i ], | X | \leq k, | Y | \leq k, | Y \backslash X | \leq 1 \right\}.
$$
Definition D.28 (Universal Dynamic Condition 2). For  $i \in [n]$ , we define  $\mathrm{opt}_i$  as Definition D.27. We say it has universal dynamic condition if for each  $i \in [n]$ , we have
$$
\operatorname {o p t} _ {i} \geq (1 - \gamma) \cdot \operatorname {o p t} _ {i - 1}.
$$
D.12 Induction Lemma for Exact Function
The goal of this section is to prove Lemma D.29
Lemma D.29 (Induction Lemma). For a fixed  $i$ , suppose the following conditions hold
- Set Condition.  $S_{i} \subseteq [i - 1]$
- Budget Condition.  $|S_{i}| \leq k$
- Value Condition.  $f_{S_{i - 1},i - 1}(S_i) \geq (1 - 1 / e) \cdot (1 - \theta)^i (1 - \gamma)^i \mathsf{opt}_i$
- Universal Dynamic Condition 1 (for exact function). (See Definition D.25)  $f_{S_i,i}(S_i) \geq (1 - \theta) \cdot f_{S_{i-1},i-1}(S_i)$
- Universal Dynamic Condition 2. (See Definition D.28)  $\mathrm{opt}_i\geq (1 - \gamma)\cdot \mathrm{opt}_{i + 1}$
- Universal Monotone Condition (for exact function). (See Definition D.23)  $f_{S_{i},i}(X) \geq f_{S_{i},i}(Y)$  for all  $Y \subset X$
Then if we construct  $S_{i+1}$  as Definition D.20, then we have
- Set Condition.  $S_{i + 1} \subseteq [i]$
- Budget Condition.  $\left| {S}_{i + 1}\right|  \leq  k$
- Value Condition.  $f_{S_{i,i}}(S_{i + 1}) \geq (1 - 1 / e) \cdot (1 - \theta)^{i + 1}(1 - \gamma)^{i + 1} \mathsf{opt}_{i + 1}$
Proof. Proof of Set Condition.
Note that  $S_{i} \subseteq [i - 1]$ , by using the way we construct  $S_{i + 1}$ , then it is obvious that  $S_{i + 1} \subseteq [i]$ .
Proof of Budget Condition.
Note that  $|S_i| \leq k$ , by using the way we construct  $S_{i + 1}$ , then it is straightforward that  $|S_{i + 1}| \leq k$ .
Proof of Value Condition.
$$
\begin{array}{l} f _ {S _ {i}, i} \left(S _ {i + 1}\right) \geq f _ {S _ {i}, i} \left(S _ {i}\right) \\ \geq (1 - \theta) \cdot f _ {S _ {i - 1}, i - 1} \left(S _ {i}\right) \\ \geq (1 - \theta) \cdot ((1 - 1 / e) (1 - \theta) ^ {i} (1 - \gamma) ^ {i} \mathrm {o p t} _ {i} - i \cdot \epsilon_ {0}) \\ \geq (1 - \theta) \cdot ((1 - 1 / e) (1 - \theta) ^ {i} (1 - \gamma) ^ {i + 1} \mathsf {o p t} _ {i + 1}) \\ \geq (1 - 1 / e) \cdot (1 - \theta) ^ {i + 1} (1 - \gamma) ^ {i + 1} \mathbf {o p t} _ {i + 1}. \\ \end{array}
$$
where the first step follows from Universal Monotone Condition, the second step follows from Universal Dynamic Condition 1, the third step follows from Value Condition, the forth step follows from Universal Dynamic Condition 2, and the last step follows from simple algebra.
D.13 Induction Lemma for Approximate Function
The goal of this section is to prove Lemma D.30
Lemma D.30 (Induction Lemma). For a fixed  $i$ , suppose the following conditions hold
- Set Condition.  $S_{i} \subseteq [i - 1]$
- Budget Condition.  $\left| {S}_{i}\right|  \leq  k$
- Value Condition.  $f_{S_{i - 1},i - 1}(S_i)\geq (1 - 1 / e)\cdot (1 - \theta)^i (1 - \gamma)^i\mathrm{opt}_i - i\cdot \epsilon_0$
- Universal Dynamic Condition 1 (for approximate function). (see Definition D.26)  $\widetilde{f}_{S_{i},i}(S_{i})\geq (1 - \theta)\cdot \widetilde{f}_{S_{i - 1},i - 1}(S_{i})$
- Universal Dynamic Condition 2.  $\mathrm{opt}_i\geq (1 - \gamma)\cdot \mathrm{opt}_{i + 1}$
- Universal Approximate Condition. (See Definition D.28)  $f_{S_i,i}(X)\geq \widetilde{f}_{S_i,i}(X) - \epsilon_0$  for all  $X$
- Universal Monotone Condition (for approximate function). (see Definition D.24)  $\widetilde{f}_{S_i,i}(X) \geq \widetilde{f}_{S_i,i}(Y)$  for all  $Y \subset X$
Then if we construct  $S_{i+1}$  as Definition D.22, then we have
- Set Condition.  $S_{i+1} \subseteq [i]$
- Budget Condition.  $\left| {S}_{i + 1}\right|  \leq  k$
- Value Condition.  $f_{S_i,i}(S_{i+1}) \geq (1 - 1/e) \cdot (1 - \theta)^{i+1}(1 - \gamma)^{i+1} \mathsf{opt}_{i+1} - (i + 1) \cdot \epsilon_0$
$$
\begin{array}{l} f _ {S _ {i}, i} (S _ {i + 1}) \geq \widetilde {f} _ {S _ {i}, i} (S _ {i + 1}) - \epsilon_ {0} \\ \geq \widetilde {f} _ {S _ {i}, i} (S _ {i}) - \epsilon_ {0} \\ \geq (1 - \theta) \cdot \tilde {f} _ {S _ {i - 1}, i - 1} \left(S _ {i}\right) - \epsilon_ {0} \\ \geq (1 - \theta) \cdot ((1 - 1 / e) (1 - \theta) ^ {i} (1 - \gamma) ^ {i} \mathbf {o p t} _ {i} - i \cdot \epsilon_ {0}) - \epsilon_ {0} \\ \geq (1 - \theta) \cdot ((1 - 1 / e) (1 - \theta) ^ {i} (1 - \gamma) ^ {i + 1} \mathbf {o p t} _ {i + 1} - i \cdot \epsilon_ {0}) - \epsilon_ {0} \\ \geq (1 - 1 / e) \cdot (1 - \theta) ^ {i + 1} (1 - \gamma) ^ {i + 1} \mathsf {o p t} _ {i + 1} - (i + 1) \cdot \epsilon_ {0}. \\ \end{array}
$$
where the first step follows from Universal Approximate Condition, the second step follows from Universal Monotone Condition, the third step follows from Universal Dynamic Condition 1, the forth step follows from Value Condition, the fifth step follows from Universal Dynamic Condition 2, and the last step follows from simple algebra.
D.14 Theoretical Result
We first give the guarantee of our full-knowledge version (without cache size limitation).
Lemma D.31 (Formal version of Lemma 3.1). Under the mild assumption, let  $k$  denote any target size. If we greedily compute the attention score based on full information, then we can find the set  $S_{i}$  such that
$$
f \left(S _ {i}\right) \geq (1 - 1 / e) \cdot (1 - \alpha) \operatorname {o p t} _ {i},
$$
where  $\alpha \in (0,1)$  are parameters.
Proof. The proof follows from using Theorem D.17, Corollary D.18, Lemma D.29 with choosing  $\theta = \gamma = \alpha / (10n)$ .
Next, we show the guarantee for our robust and approximate greedy eviction policy algorithm (Algorithm 2).
Theorem D.32 (Formal version of Theorem 4.4). Under the mild assumption, let  $k$  denote the budget of space limitation. If for each token, we greedily compute the attention score based on top- $k$  choice, then we can show the set  $\widetilde{S}_i$  we generate each for token  $i \in [n]$  satisfy that
$$
f (\widetilde {S} _ {i}) \geq (1 - 1 / e) \cdot (1 - \alpha) \mathrm {o p t} _ {i} - \beta ,
$$
where  $\alpha \in (0,1),\beta >0$  are parameters.
Proof. The proof follows from using Theorem D.17, Corollary D.18, Lemma D.30 with choosing  $\epsilon_0 = \beta /(10n)$  and  $\theta = \gamma = \alpha /(10n)$ .
D.15 Extended Related Work for Theoretical Attention Problems
The static attention computation is asking the following question that given  $Q, K, V \in \mathbb{R}^{n \times d}$ , the goal is to  $D^{-1} \exp(QK^{\top})V$  where  $D = \mathrm{diag}(\exp(QK^{\top})\mathbf{1}_n)$ . [98] studied the static attention computation from both algorithm and hardness. On the positive, they provide an almost linear time algorithm to approximately compute the attention matrix. On the negative side, assuming a strong exponential time hypothesis (SETH), they prove a hardness result. Their hardness result is, unless SETH fails, there is no algorithm that runs in truly subquadratic time to approximately compute the attention matrix. Further, [100] considers the dynamic of attention computation problem. They also provide both algorithmic results and hardness results. In the work of [103], they consider the sparsification of attention matrix construction. In particular, they assume that situation that  $d \gg n$ , and show how to sparsify the columns of matrix  $Q$ . [103] provides two algorithms, one is a randomized algorithm, and the other is a deterministic algorithm. Differential privacy is a famous and textbook topic in graduate school, recently the work of [101] shows how to give a differentially private algorithm for computing the attention matrix. For a given  $A \in \mathbb{R}^{n \times d}$  and vector  $b \in \mathbb{R}^n$ , [104] formulates and studies exponential regression  $\min_x \| \exp(Ax) - b \|_2$ . Then [105] considers the normalization factor in exponential regression and defines the softmax regression problem  $\min_x \| \langle \exp(Ax), \mathbf{1}_n \rangle^{-1} \exp(Ax) - b \|_2$ . [107] moves the scaling factor from  $\exp(Ax)$  to  $b$  and defines a rescaled softmax regression problem  $\min_x \| \exp(Ax) - \langle \exp(Ax), \mathbf{1}_n \rangle \cdot b \|_2$ .
D.16 Sparsity Preserving
Recall that in Figure 2, we observe that even when trained densely, the attention matrices of LLMs are over  $95\%$  sparse at inference time. Only  $5\%$  of the KV cache is sufficient for decoding the same output token at each generation step. Here, we provide some formal formulations for sparsity.
Definition D.33. Suppose the following conditions
- Let  $S_0 \subset [m]$ .
- Let  $k = |S_0|$ .
- Let  $\tau \in (0,1)$  denote a threshold for truncating the value.
- Let  $\alpha \in (0,1)$  denote a fraction of mass (larger than  $\tau$ ) outside  $S_0$ .
- Let mapping  $\mathcal{D}:\mathbb{R}^d\to \mathbb{R}_{\geq 0}^m$
- For each  $x \in \mathbb{R}^d$ ,  $\mathcal{D}(x) \in \mathbb{R}^m$  is a vector that has length  $m$ .
We say the distribution  $\mathcal{D}$  is  $(\alpha, \tau, k)$ -good if the following conditions hold
- For all  $x \in \mathbb{R}^d$ ,  $S_0 \subset \operatorname{supp}_{\tau}(\mathcal{D}(x))$
- For all  $x \in \mathbb{R}^d$ ,  $|\operatorname{supp}_{\tau}(\mathcal{D}(x)) \backslash S_0| \leq \alpha \cdot k$
Claim D.34. Suppose we sample  $n$  points  $\{x_{1}, x_{2}, \dots, x_{n}\} \subset \mathbb{R}^{d}$  from  $(\alpha, \tau, k)$ -good distribution uniformly at random, then we have
-  $S_0 \subseteq \cap_{i \in [n]} \operatorname{supp}_{\tau}(x_i)$
-  $\left| {\left( {{\cup }_{i \in  \left\lbrack  n\right\rbrack  }\operatorname{supp}_{\tau }\left( \mathcal{D}\left( x\right) \right) }\right)  \backslash  {S}_{0}}\right|  \leq  {\alpha kn}$
Proof. Since for all  $i \in [n]$ , we have  $S_0 \subseteq \operatorname{supp}_{\tau}(\mathcal{D}(x_i))$ , thus
$$
S _ {0} \subseteq \cap_ {i \in [ n ]} \operatorname {s u p p} _ {\tau} (x _ {i}).
$$
Therefore we proved the first property.
We know that for all  $i\in [n]$  , we have  $\left|\operatorname {supp}_{\tau}(\mathcal{D}(x))\backslash S_0\right|\leq \alpha kn$  . Thus
$$
\left| \left(\cup_ {i \in [ n ]} \operatorname {s u p p} _ {\tau} (\mathcal {D} (x _ {i}))\right) \backslash S _ {0} \right| \leq \sum_ {i = 1} ^ {n} \left| \operatorname {s u p p} _ {\tau} (\mathcal {D} (x _ {i}))\right) \backslash S _ {0} | \leq n \cdot \alpha k
$$
Therefore, we finish the proof for the second property.
D.17 Definition of Loss Function
In this section, we follow the theoretical softmax regression literature [105] and define a number of functions to make the calculations of gradient and Hessian convenient. We also proposed a new penalty term  $(\ell_{1}$  type sparsity penalty, see Definition D.41) into the final loss function, which is not studied in previous work [104, 105, 107, 110, 132, 133]. We first provide some function definitions.
Definition D.35 (Function  $u$ , [105]). Given matrix  $A \in \mathbb{R}^{n \times d}$ , let function  $u: \mathbb{R}^d \to \mathbb{R}^n$  be defined as follows
$$
u (x) := \exp (A x)
$$
Definition D.36 (Function  $\alpha$ , see Definition 5.4 in [105] as an example). We define  $u(x)$  as Definition D.35. Then we define  $\alpha: \mathbb{R}^d \to \mathbb{R}$  as follows
$$
\alpha (x) := \langle u (x), \mathbf {1} _ {n} \rangle
$$
Definition D.37 (Function  $f$ , see Definition 5.1 in [105] as an example). Provided that the following conditions are true
- We define  $u(x)$  as Definition D.35.
- We define  $\alpha(x)$  as Definition D.36
Let function  $f: \mathbb{R}^d \to \mathbb{R}^n$  be defined as follows
$$
f (x) := \alpha (x) ^ {- 1} u (x).
$$
Definition D.38 (Function  $c$ , see Definition 5.5 in [105] as an example). Provided that the following conditions are true
- Given a vector  $b \in  {\mathbb{R}}^{n}$  .
- Let  $f(x)$  be defined as Definition D.38.
Then, let function  $c:\mathbb{R}^d\to \mathbb{R}^n$  defined as follows
$c(x)\coloneqq f(x) - b.$
Definition D.39 (Loss function  $L_{\mathrm{exp}}$ , see Definition 5.3 in [105] as an example). We define  $L_{\mathrm{exp}}: \mathbb{R}^d \to \mathbb{R}$
$$
L _ {\exp} (x) := 0. 5 \cdot \| c (x) \| _ {2} ^ {2}.
$$
Definition D.40 (Loss function  $L_{\mathrm{reg}}$ ). Given  $A \in \mathbb{R}^{n \times d}$ .
Let function  $L_{\mathrm{reg}}:\mathbb{R}^d\to \mathbb{R}$  be defined as follows
$$
L _ {\text {r e g}} (x) := 0. 5 \cdot \| \operatorname {d i a g} (w) A x \| _ {2} ^ {2}
$$
We define a novel penalty function
Definition D.41 (Implicitly controlling the sparsity). Given  $A \in \mathbb{R}^{n \times d}$
$$
L _ {\text {s p a r s e}} (x) := \left\| \exp (A x) \right\| _ {1}.
$$
Then it is obvious that we have
Claim D.42. Given  $A \in \mathbb{R}^{n \times d}$ . Let  $u$  be defined as Definition D.35. Let  $\alpha$  be defined as Definition D.36.
-  $L_{\text{sparse}}(x) = \langle \exp(Ax), \mathbf{1}_n \rangle$
-  $L_{\mathrm{sparse}}(x) = \langle u(x), \mathbf{1}_n \rangle$
-  $L_{\text{sparse}}(x) = \alpha(x)$
Proof. The proof is trivially following from the definition of  $u(x)$  (see Definition D.35) and  $\alpha(x)$  (see Definition D.36).
The final loss function can be defined as follows. Intuitively, we can write attention  $D^{-1}\exp (QK^{\top})$  into  $n$  subproblems where each subproblem can be viewed as one softmax problem.
Definition D.43. If the following conditions hold
- We define  $L_{\mathrm{exp}}$  as Definition D.39.
- We define  $L_{\mathrm{reg}}$  as Definition D.40.
- We define  $L_{\text{sparse}}$  as Definition D.41.
Then we define  $L$  function
$$
L (x) := L _ {\exp} (x) + L _ {\text {s p a r s e}} (x) + L _ {\text {r e g}} (x).
$$
D.18 Gradient
Next, we show the gradient of  $L_{\mathrm{exp}}$ .
Lemma D.44 (Gradient, Lemma 5.6 in [105]). Provided that the following conditions are true
- Given matrix  $A \in \mathbb{R}^{n \times d}$  and a vector  $b \in \mathbb{R}^n$ .
- Let  $A_{*,i} \in \mathbb{R}^n$  denote the  $i$ -th column of matrix  $A$ , for every  $i \in [d]$ .
- We define  $\alpha(x)$  as Definition D.36.
- We define  $f(x)$  as Definition D.37.
- We define  $c(x)$  as Definition D.38.
- We define  $L_{\mathrm{exp}}(x)$  as Definition D.39.
- Let  $\circ$  denote hadamard product.
For every  $i\in [d]$  ,we have
$$
\frac {\mathrm {d} \exp (A x)}{\mathrm {d} x _ {i}} = \exp (A x) \circ A _ {*, i}
$$
$$
\frac {\mathrm {d} \left\langle \exp (A x) , \mathbf {1} _ {n} \right\rangle}{\mathrm {d} x _ {i}} = \left\langle \exp (A x), A _ {*, i} \right\rangle
$$
$$
\frac {\mathrm {d} \alpha (x) ^ {- 1}}{\mathrm {d} x _ {i}} = - \alpha (x) ^ {- 1} \cdot \langle f (x), A _ {*, i} \rangle
$$
$$
\frac {\mathrm {d} f (x)}{\mathrm {d} x _ {i}} = \frac {\mathrm {d} c (x)}{\mathrm {d} x _ {i}} = - \langle f (x), A _ {*, i} \rangle \cdot f (x) + f (x) \circ A _ {*, i}
$$
$$
\frac {\mathrm {d} L _ {\exp} (x)}{\mathrm {d} x _ {i}} = A _ {*, i} ^ {\top} \cdot (f (x) (f (x) - b) ^ {\top} f (x) + \operatorname {d i a g} (f (x)) (f (x) - b))
$$
D.19 Hessian
Here, we compute the Hessian for several functions.
Lemma D.45 (Hessian of  $u(x)$ , Lemma 5.9 in [105]). Provided that the following conditions are true
- Given a matrix  $A \in \mathbb{R}^{n \times d}$ .
- For every  $i \in [d]$ , let  $A_{*,i} \in \mathbb{R}^n$  denote the  $i$ -th column of matrix  $A$ .
- Let  $\circ$  denote hadamard product.
Then, we have, for each  $i \in [d]$
$$
\frac {\mathrm {d} ^ {2} \exp (A x)}{\mathrm {d} x _ {i} ^ {2}} = A _ {*, i} \circ u (x) \circ A _ {*, i}
$$
$$
\frac {\mathrm {d} ^ {2} \exp (A x)}{\mathrm {d} x _ {i} \mathrm {d} x _ {j}} = A _ {*, j} \circ u (x) \circ A _ {*, i}
$$
Lemma D.46 (Lemma 5.10 in [105]). Provided that the following conditions are true
- We define  $\alpha(x)$  as Definition D.36.
- For every  $i \in [d]$ , let  $A_{*,i} \in \mathbb{R}^n$  denote the  $i$ -th column of matrix  $A$ .
- Let  $\circ$  denote hadamard product.
$$
\frac {\mathrm {d} ^ {2} \alpha (x)}{\mathrm {d} x _ {i} ^ {2}} = \langle u (x), A _ {*, i} \circ A _ {*, i} \rangle
$$
$$
\frac {\mathrm {d} ^ {2} \alpha (x)}{\mathrm {d} x _ {i} \mathrm {d} x _ {j}} = \langle u (x), A _ {*, i} \circ A _ {*, j} \rangle
$$
$$
\frac {\mathrm {d} ^ {2} \alpha (x)}{\mathrm {d} x ^ {2}} = A ^ {\top} \operatorname {d i a g} (u (x)) A
$$
D.20 Hessian is Positive Definite
It is well known that in literature [104, 105, 107], the Hessian  $H$  of loss function can be written as  $A^{\top}(B(x) + W^{2})A$  for some matrix function  $B(x)\in \mathbb{R}^{n\times n}$  (for example see explanation in Section 5.10 in [105]). In this section, we show that Hessian is positive definite.
Lemma D.47. If the following conditions hold
- Given matrix  $A \in \mathbb{R}^{n \times d}$ .
- We define  $L_{\mathrm{sparse}}(x)$  as Definition D.41.
- We define  $L_{\mathrm{sparse}}(x)$  Definition D.40.
- We define  $L_{\mathrm{sparse}}(x)$  Definition D.39.
- Let  $L(x) = L_{\exp}(x) + L_{\mathrm{sparse}}(x) + L_{\mathrm{reg}}(x)$ .
- Let  $A^\top (B(x) + W^2)A$  be the Hessian of  $L(x)$ .
- Let  $W = \operatorname{diag}(w) \in \mathbb{R}^{n \times n}$ . Let  $W^2 \in \mathbb{R}^{n \times n}$  denote the matrix that  $i$ -th diagonal entry is  $w_{i,i}^2$ .
- Let  $\sigma_{\min}(A)$  denote the minimum singular value of  $A$ .
- Let  $l > 0$  denote a scalar.
- Part 1. If all  $i \in [n]$ ,  $w_i^2 \geq 20 + l / \sigma_{\min}(A)^2$ , then
$$
\frac {\mathrm {d} ^ {2} L}{\mathrm {d} x ^ {2}} \succeq l \cdot I _ {d}
$$
- Part 2 If all  $i \in [n]$ ,  $w_i^2 \geq 200 \cdot \exp(R^2) + l / \sigma_{\min}(A)^2$ , then
$$
(1 - 1 / 1 0) \cdot (B (x) + W ^ {2}) \preceq W ^ {2} \preceq (1 + 1 / 1 0) \cdot (B (x) + W ^ {2})
$$
Proof. The entire proof framework follows from [104, 105, 107], in the next few paragraphs, we mainly explain the difference.
The  $B(x)$  based on  $L_{\mathrm{sparse}}$  is  $\operatorname {diag}(u(x))$  . Note that it is obvious that
$$
\operatorname {d i a g} (u (x)) \succeq 0.
$$
From the upper bound size, we know that
$$
\begin{array}{l} \operatorname {d i a g} (u (x)) \preceq \| u (x) \| _ {\infty} \cdot I _ {n} \\ \succeq \exp (R ^ {2}) \\ \end{array}
$$
where the last step follows from Proof of Part 0 in Lemma 7.2 in [105].
To prove Part 1, following from [104, 105], we only use the lower bound of  $\mathrm{diag}(u(x))$ . By putting things together, we get our results.
To prove Part 2, we follow from [107] and use both the upper bound and lower bound of  $\mathrm{diag}(u(x))$ .
D.21 Hessian is Lipschitz
In this section, we show Hessian is Lipschitz.
Lemma D.48. If the following conditions hold
- Let  $H(x) = A^{\top}(B(x) + W^{2})A$  denote the Hessian of  $L$ .
-  $\| H(x) - H(y)\| \leq n^2\exp (40R^2)\| x - y\| _2$
Note that the  $B(x)$  based on  $L_{\mathrm{exp}} + L_{\mathrm{reg}}$  have been proved by [105]
We only need to prove  $B(x)$  based on  $L_{\mathrm{sparse}}$  and add them together.
Note that  $B(x)$  based on  $L_{\mathrm{sparse}}$  is in fact  $\operatorname {diag}(u(x))$
Using Lemma 7.2 in [105], we know that
$$
\begin{array}{l} \left\| \operatorname {d i a g} (u (x)) - \operatorname {d i a g} (u (y)) \right\| \leq \| u (x) - u (y) \| _ {2} \\ \leq 2 \sqrt {n} R \exp (R ^ {2}) \| x - y \| _ {2} \\ \end{array}
$$
where the last step follows from Part 1 in Lemma 7.2 in [105].
Thus, putting things together, we complete the proof.
D.22 Greedy Type Algorithm
In this section, we propose a greedy-type algorithm (based on the approximate Newton method) to solve the optimization problem.
Algorithm 4 A greedy type algorithm.
1: procedure OURITERATIVEMETHOD  $(A\in \mathbb{R}^{n\times d},b\in \mathbb{R}^n,w\in \mathbb{R}^n,\epsilon ,\delta)$    
2: Initialize  $x_0$    
3:  $T\gets \log (\| x_0 - x^*\| _2 / \epsilon)$  Let  $T$  denote the number of iterations.   
4: for  $t = 0\rightarrow T$  do   
5:  $D\gets B_{\mathrm{diag}}(xt) + \mathrm{diag}(w\circ w)$    
6:  $\widetilde{D}\gets \mathrm{SUBSAMPLE}(D,A,\epsilon_1 = \Theta (1),\delta_1 = \delta /T)$    
7: Compute gradient  $g$  exactly   
8: Get the approximate Hessian  $\tilde{H}$  by computing  $A^{\top}\tilde{D} A$    
9: Update  $x_{t + 1}$  by using the Newton step  $x_{t} + \tilde{H}^{-1}g$    
10: end for   
11:  $\widetilde{x}\gets x_{T + 1}$    
12: return  $\widetilde{x}$    
13: end procedure
Theorem D.49. Given matrix  $A \in \mathbb{R}^{n \times d}$ ,  $b \in \mathbb{R}^n$ , and  $w \in \mathbb{R}^n$ .
- We use  $x^{*}$  to denote the optimal solution of
$$
\min  _ {x \in \mathbb {R} ^ {d}} L _ {\text {e x p}} + L _ {\text {s p a r s e}} + L _ {\text {r e g}}
$$
$g(x^{*}) = \mathbf{0}_{d}$  where  $g$  denotes the gradient function.
-  $\| x^{*}\|_{2}\leq R$