SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Devvrit *†
devvrit.03@gmail.com
Sai Surya Duvvuri*
subramanyamdvss@gmail.com
Rohan Anil
rohananil@google.com
Vineet Gupta
vineet@google.com
Cho-Jui Hsieh
CS Department, UCLA & Google
chohsieh@cs.ucla.edu
Inderjit Dhillon
isd@google.com
Second-order methods hold significant promise for enhancing the convergence of deep neural network training; however, their large memory and computational demands have limited their practicality. Thus there is a need for scalable second-order methods that can efficiently train large models. In this paper, we introduce the Sparsified Online Newton (SONew) method, a memory-efficient second-order algorithm that yields a sparsified yet effective preconditioner. The algorithm emerges from a novel use of the LogDet matrix divergence measure; we combine it with sparsity constraints to minimize regret in the online convex optimization framework. Empirically, we test our method on large scale benchmarks of up to 1B parameters. We achieve up to  $30\%$  faster convergence,  $3.4\%$  relative improvement in validation performance, and  $80\%$  relative improvement in training loss, in comparison to memory efficient optimizers including first order methods. Powering the method is a surprising fact - imposing structured sparsity patterns, like tridiagonal and banded structure, requires little to no overhead, making it as efficient and parallelizable as first-order methods. In wall-clock time, tridiagonal SONew is only about  $3\%$  slower per step than first-order methods but gives overall gains due to much faster convergence. In contrast, one of the state-of-the-art (SOTA) memory-intensive second-order methods, Shampoo, is unable to scale to large benchmarks. Additionally, while Shampoo necessitates significant engineering efforts to scale to large benchmarks, SONew offers a more straightforward implementation, increasing its practical appeal. SONew code is available at: https://github.com/devvrit/SONew
Stochastic first order methods which use the negative gradient direction to update parameters have become the standard for training deep neural networks (DNNs). Gradient-based preconditioning involves finding an update direction, by multiplying the gradient with a preconditioner matrix carefully chosen from gradients observed in previous iterations, to improve convergence. (Full
* equal contribution, † Work done while at Google
matrix) Adagrad [15], online Newton method [25] and natural gradient descent [3] use a full-matrix preconditioner, but computing and storing the full matrix is infeasible when there are millions of parameters. Thus, diagonal versions such as diagonal Adagrad, Adam [33], and RMSprop [28] are now widely used to train DNNs due to their scalability.
Several higher-order methods have previously been applied to deep learning ([24, 5, 23, 38]). All these methods use Kronecker product factorizations that reduce computational and storage costs to make them feasible for training neural networks. However, to precondition a  $d_{1} \times d_{2}$  parameter matrix, these methods require matrix inverse operations, which take  $\mathcal{O}(d_1^3 + d_2^3)$  time and  $\mathcal{O}(d_1^2 + d_2^2)$  space. In comparison, first-order methods use  $\mathcal{O}(d_1 d_2)$  time and memory, which is linear in the number of parameters. For instance, when  $d_{1} = k d_{2}$ , the memory used by Shampoo,  $d_{1}^{2} + d_{2}^{2}$  floating point numbers is  $\mathcal{O}(k)$  times the number of parameters, which could be arbitrarily large depending on  $k$ . This calls for further research in developing efficient second-order optimization techniques to train DNNs with memory and time complexity linear in the number of parameters.
In this paper, we present a novel Sparsified Online Newton (SONew) method, which only requires linear time and space complexity, to train large-scale DNNs. We derive the algorithm through two steps, classical regret analysis followed by a sparsification step. In more detail, regret analysis when using a preconditioner reveals that the error is bounded by two terms, the first depends on the change in the preconditioning matrix, while the second depends on the generalized gradient norm (see Section 3 for more details). We take a novel approach of minimizing the second term while regularizing two successive preconditioners to be close in the LogDet matrix divergence measure [34] (see Section 3 for the intuition behind choosing LogDet divergence). This analysis naturally yields us an Online Newton method [25]. To make it computationally efficient, we further sparsify the preconditioner by finding a sparse approximation that is close in LogDet divergence. Thus we are consistent in using the same measure (LogDet divergence) in both the regularization and sparsification steps. This gives us our SONew method, which achieves linear complexity by leveraging structured sparsity patterns, such as tridiagonal and banded, in the preconditioner. This is unlike most existing online Newton methods that require quadratic space and cubic time complexity. By making each step linear time, the SONew method can be applied to train modern DNNs as efficiently as first order methods. Further, our method is embarrassingly parallelizable thus making negligible the overhead of computing the preconditioner. We also show that introducing sparsity allows us to reduce the condition number of the problem dynamically to improve numerical stability.
We strengthen the relationship between sparse LogDet divergence minimization and online convex optimization by establishing an optimal  $\mathcal{O}(\sqrt{T})$  regret upper bound for tridiagonal sparsity pattern. In our experiments on an MLP Autoencoder and Graph Neural Network (GNN), we found that our method outperformed first-order methods in terms of training loss within the same training time, while Shampoo (second-order method) takes significantly longer. In our experiments on Vision Transformers on Imagenet and GNN on OGBG-molpcba, we achieve a target validation performance using  $10\%$  and  $30\%$  fewer iterations respectively compared to Adam, the SOTA optimizer for both benchmarks. Furthermore, using the same number of iterations as Adam we observe  $0.7\%$  and  $3.4\%$  relative improvement for ViT and GNN respectively in validation performance. From an optimization point of view, SONew achieves  $9\%$  and  $80\%$  better relative training loss for ViT and GNN respectively. It is worth noting that Shampoo statistics required  $\sim 7 \times \# \text{params}$  for ViT whereas tridiag-SONew uses only  $2 \times \# \text{params}$  for its statistics. We also test another recently proposed memory efficient second order optimizer, rfdSON [37], but found its performance suboptimal to the best performing first order method. Owing to SONew's scalability, we train a Large Language Model (LLM) with 1 billion parameters and compare it with AdaFactor [45], a popularly used first order optimizer to train LLMs [11]. SONew achieves the same performance as AdaFactor using  $26\%$  fewer steps, resulting in a  $1.35 \times$  faster training. When using the same number of steps, SONew obtained a  $1.7\%$  relative better train loss. In terms of implementation, SONew is just a few lines of code (Equation (13)) without complex engineering challenges, rendering it even more useful and practical.
The inner product between matrices is defined as  $\langle A, B \rangle = \operatorname{Tr}(A^T B)$ , where  $\operatorname{Tr}(\cdot)$  denotes the matrix trace. The Frobenius norm of a matrix  $A$  is  $\|A\|_F = \sqrt{\operatorname{Tr}(A^TA)}$ , while its spectral norm is  $\|A\|_2 = \max_x \|Ax\|_2 / \|x\|_2$ . We use  $I_n \in \mathbb{R}^{n \times n}$  to denote an identity matrix. We use  $S_n$ ,  $S_n^{++}$  to denote the set of symmetric, and positive definite matrices respectively. The generalized norm of a
vector  $x \in \mathbb{R}^n$  with respect to matrix  $A \in S_n^{++}$  is defined as  $\| x \|_A = \sqrt{x^TAx}$ . We use  $\operatorname*{det}(A)$  to denote the determinant of matrix  $A$ , and  $\mathrm{diag}(A)$  to denote the diagonal matrix with  $\mathrm{diag}(A)_{ii} = A_{ii}$ . We use  $\mathcal{G}$  and  $\tilde{\mathcal{G}}$  to denote a graph and its sub-graph with a vertex set  $[n] = \{1, \ldots, n\}$ . Let  $E_{\mathcal{G}}$  denote the set of edges in graph  $\mathcal{G}$ , and  $\mathrm{neig}_{\mathcal{G}}(i)$  denote neighbours of vertex  $i$  in graph  $\mathcal{G}$ . A sparse symmetric matrix  $A \in \mathbb{R}^{n \times n}$  follows a sparsity structure graph  $\mathcal{G}$  if  $A_{i,j} = 0 \forall (i,j) \notin E_{\mathcal{G}}$ . Note that the set of all such matrices forms a linear subspace. We use  $S_n(\mathcal{G})^{++}$  to denote the set of positive definite matrices with sparsity structure given by graph  $\mathcal{G}$ , i.e., if  $X \in S_n(\mathcal{G})^{++}$ , then  $X_{ij} = 0 \forall (i,j) \notin E(\mathcal{G})$ .  $S_n(\mathcal{G})^{++}$  is an open convex set. Given an index set  $I = \{i_1, i_2, \ldots, i_n\}$ , we use  $A_{II}$  to denote the corresponding principal sub-matrix of  $A$ .
2.1 LogDet matrix divergence
Let  $\phi : S_n^{++} \to \mathbb{R}$  be a strictly convex, differentiable function. The Bregman matrix divergence between  $X, Y \in S_n^{++}$  is defined as [8, 34]:  $\mathrm{D}_{\phi}(X,Y) = \phi(X) - \phi(Y) - \mathrm{Tr}(\nabla \phi(Y)^T(X - Y))$ . Since  $\phi$  is convex,  $\mathrm{D}_{\phi}(X,Y) \geq 0$  for all  $X,Y \in S_n^{++}$ . For example if  $\phi(X) = \|X\|_F^2$ , the corresponding Bregman divergence  $\mathrm{D}_{\phi}(X,Y) = \|X - Y\|_F^2$  is the squared Frobenius distance. In this paper, we extensively use the convex function  $\phi(X) = -\log \det(X)$ ; the corresponding divergence measure  $\mathrm{D}_{\ell\mathrm{d}}(X,Y)$  is called the LogDet matrix divergence:
$$
\mathrm {D} _ {\ell \mathrm {d}} (X, Y) = - \log \det  \left(X Y ^ {- 1}\right) + \operatorname {T r} \left(X Y ^ {- 1}\right) - n. \tag {1}
$$
The LogDet divergence is scale invariant to invertible matrices  $A$ , i.e.  $\mathrm{D}_{\ell \mathrm{d}}(A^T X A, A^T Y A) = \mathrm{D}_{\ell \mathrm{d}}(X, Y)$ . LogDet divergence can be written in terms of eigendecompositions of  $X = V\Sigma V^T$  and  $Y = U\Theta U^T$  [34]:
$$
\mathrm {D} _ {\ell \mathrm {d}} (X, Y) = \sum_ {i} \sum_ {j} \left(v _ {i} ^ {T} u _ {j}\right) ^ {2} \left(\sigma_ {i} / \theta_ {j} - \log \left(\sigma_ {i} / \theta_ {j}\right) - 1\right). \tag {2}
$$
These two properties are later used in Section 3 to highlight the significance of LogDet divergence in our algorithm.
3 SONew: Sparsified Online Newton Method
We now present our proposed algorithm SONew.
3.1 Regret minimization via LogDet divergence
We set up our problem under the online convex optimization framework (OCO) [43, 26], where at each round the learner makes a prediction  $w_{t}$  and receives a convex loss  $f_{t}(w_{t})$  and gradient  $g_{t} = \nabla f_{t}(w_{t})$  as feedback. The goal of the learner is to reduce regret  $R_{T}$  by predicting  $w_{t}$  so that a low aggregate loss  $\sum_{t=1}^{T} f_{t}(w_{t})$  is achieved compared to the best possible,  $w^{*} = \arg \min_{w} \sum_{t=1}^{T} f_{t}(w)$ . Formally, regret is given by
$$
R _ {T} (w _ {1}, \dots , w _ {T}) = \sum_ {t = 1} ^ {T} f _ {t} (w _ {t}) - \sum_ {t = 1} ^ {T} f _ {t} (w ^ {*}).
$$
Using [10],  $R$  regret in online setting yields  $R / T$  convergence rate in the stochastic setting. To upper bound this regret, we proceed as in [26] by analyzing the error in the iterates for the update  $w_{t + 1} \coloneqq w_t - \eta X_t g_t$ , where  $X_t \in \mathbb{R}^{n \times n}$ . Then  $\| w_{t + 1} - w^* \|_{X_t^{-1}}^2 = \| w_t - \eta X_t g_t - w^* \|_{X_t^{-1}}^2 = \| w_t - w^* \|_{X_t^{-1}}^2 + \eta^2 g_t^T X_t g_t - 2\eta (w_t - w^*)^T g_t$ . The convexity of  $f_t$  implies that  $f_t(w_t) - f_t(w^*) \leq (w_t - w^*)^T g_t$  leading to  $f_t(w_t) - f_t(w^*) \leq \frac{1}{2\eta} (\| w_t - w^* \|_{X_t^{-1}}^2 - \| w_{t + 1} - w^* \|_{X_t^{-1}}^2 + \eta^2 g_t^T X_t g_t)$ . Summing over all  $t \in [T]$  and rearranging reveals the following upper bound on overall regret:
$$
R _ {T} \leq \frac {1}{2 \eta} \| w _ {1} - w ^ {*} \| _ {X _ {1} ^ {- 1}} ^ {2} + \frac {\eta}{2} \sum_ {t = 1} ^ {T} g _ {t} ^ {T} X _ {t} g _ {t} + \frac {1}{2 \eta} \sum_ {t = 2} ^ {T} \left(w _ {t} - w ^ {*}\right) ^ {T} \left(X _ {t} ^ {- 1} - X _ {t - 1} ^ {- 1}\right) \left(w _ {t} - w ^ {*}\right). \tag {3}
$$
Since  $w^{*}$  is unknown, finding  $X_{t}$  which minimizes (3) is infeasible. So to minimize regret, we attempt to minimize the second term in (3) while regularizing  $X_{t}^{-1}$  to be "close" to  $X_{t - 1}^{-1}$ . The nearness measure we choose is the LogDet matrix divergence, thus leading to the following objective
$$
X _ {t} = \underset {X \in S _ {n} ^ {+ +}} {\arg \min } g _ {t} ^ {T} X g _ {t}, \text {s u c h t h a t} \mathrm {D} _ {\ell \mathrm {d}} (X, X _ {t - 1}) \leq c _ {t}, \tag {4}
$$
where  $\mathrm{D}_{\ell \mathrm{d}}$  is as in (1). Why do we use the LogDet divergence? From (2), due to the term  $\lambda_{i} / \theta_{j}$ ,  $\mathrm{D}_{\ell \mathrm{d}}(X,X_{t - 1})$  prioritizes matching the smaller eigenvalues of  $X_{t - 1}$  with those of  $X$ , i.e., matching the larger eigenvalues of  $X_{t - 1}^{-1}$  and  $X^{-1}$ . As a consequence, LogDet divergence regularizes  $X$  by matching up its large eigenvalues with those of  $X_{t - 1}$ . For example if smallest and largest eigenvalue of  $X_{t - 1}$  are  $\theta_{n}$  and  $\theta_{1}$ , then for an eigenvalue  $\sigma$  of  $X$ , when  $\sigma >\theta_{n}$ ,  $\theta_{1}$ , the penalty from (2) for  $\theta_{n}$  is higher than for  $\theta_{1}$ ,  $(\sigma /\theta_{n} - \log (\sigma /\theta_{n}) - 1) > (\sigma /\theta_{1} - \log (\sigma /\theta_{1}) - 1)$ . This intuition leads us to formulate (4) as our objective. We recall that there is precedence of using the LogDet divergence in the optimization literature; indeed the celebrated BFGS algorithm [9, 17, 22, 44] can be shown to be the unique solution obtained when the LogDet divergence between successive preconditioners, subject to a secant constraint, is minimized (as shown in the 4-page paper by [18]).
The optimization problem in (4) is convex in  $X$  since the LogDet divergence is convex in its first argument. The Lagrangian  $\mathcal{L}(X, \lambda_t) = g_t^T X g_t + \lambda_t (\mathrm{D}_{\ell \mathrm{d}}(X, X_{t-1}) - c_t) = \mathrm{Tr}(X g_t g_t^T) + \lambda_t (-\log \det(X X_{t-1}^{-1}) + \mathrm{Tr}(X X_{t-1}^{-1}) - n)) - \lambda_t c_t$ . Setting  $\nabla \mathcal{L}(X, \lambda_t) = 0$ , and using the fact that  $\nabla \log \det(X) = X^{-1}$  we get the following update rule:
$$
X _ {t} ^ {- 1} = X _ {t - 1} ^ {- 1} + g _ {t} g _ {t} ^ {T} / \lambda_ {t}. \tag {5}
$$
We emphasize that the update rule (5) arises naturally from our novel use of LogDet divergence to minimize the regret. Moreover, Equation (5) can be seen as a general update rule applicable to numerous existing optimizers. For example, setting  $c_{t} = 0$  (equivalently  $\lambda_t = \infty$ )  $\forall t \in [n]$  in (4) results in no change to the preconditioner in any round. In this case, with  $X_0 = I_n$ , we get online gradient descent [54]. On the other hand, setting  $\lambda_{t} = 1$  gives the update rule of the online Newton method [25]. Our update rule differs from (full-matrix) Adagrad [15] which has  $X_{t}^{-2} = X_{t - 1}^{-2} + g_{t}g_{t}^{T}$ .
Maintaining and updating  $X_{t}$  as in (5) is possible by using Sherman-Morrison formula but requires  $\mathcal{O}(n^2)$  storage and time complexity. This becomes impractical when  $n$  is in the order of millions which is typically the case in DNNs.
3.2 Sparsifying the Preconditioner
To minimize the memory needed for maintaining and updating  $X_{t}$  using (5), we adopt the strategy of sparsifying the preconditioner. For existing optimizers such as (full-matrix) Adagrad or the Online Newton method, it is unclear how to sparsify a given preconditioner. Specifically, there is no intuitive approach to assessing the quality of a sparse preconditioner compared to a full-matrix preconditioner. However, since our update rule (5) originates from using LogDet divergence in the regret bound analysis, it gives us a natural metric to measure the quality of a sparse preconditioner. Let's consider the following problem: find a sparse positive definite  $X$  with  $\| X\|_0 \leq \alpha n$ ,  $\alpha > 1$ , such that the objective  $\mathrm{D}_{\ell \mathrm{d}}(X,(X_{t - 1}^{-1} + g_t g_t^T / \lambda_t)^{-1})$  is minimized. Essentially, this problem imposes a sparsity constraint while requiring the sparse preconditioner to remain close to the full-matrix preconditioner in terms of LogDet divergence.
Due to the  $L_0$ -norm constraint, this is a non-convex problem, which makes it difficult to solve exactly. Since  $L_{1}$ -norm serves as a convex relaxation for the  $L_0$  norm, we could use it instead, resulting in the following optimization problem also known as graphical lasso estimator [19]:
$$
\min  _ {X \in S _ {n} ^ {+ +}} \mathrm {D} _ {\ell \mathrm {d}} \left(X, \left(X _ {t - 1} ^ {- 1} + g _ {t} g _ {t} ^ {T} / \lambda_ {t}\right) ^ {- 1}\right) + \gamma \| X \| _ {1}.
$$
However, the time taken to solve the above problem, even with the current best methods [7, 29, 16, 53], can still be too large (as these methods take several minutes for a matrix of size million), making it impractical to embed in DNN training.
In this paper, we take a different direction where we use fixed sparsity pattern constraints, specified by a fixed undirected graph  $\mathcal{G}$ . To sparsify the solution in (5), we formulate the subproblem
$$
X _ {t} = \underset {X \in S _ {n} (\mathcal {G}) ^ {+ +}} {\arg \min } \mathrm {D} _ {\ell \mathrm {d}} \left(X, \left(X _ {t - 1} ^ {- 1} + g _ {t} g _ {t} ^ {T} / \lambda_ {t}\right) ^ {- 1}\right), \tag {6}
$$
where  $S_{n}(\mathcal{G})^{++}$  denotes the set of positive definite matrices with the fixed sparsity pattern corresponding to the adjacency matrix of graph  $\mathcal{G}$ . Note that both steps (4) and (6) use the same LogDet measure.
Owing to the structure of LogDet divergence, (6) can be surprisingly solved in  $\mathcal{O}(n)$  and easily parallelizable, for certain sparsity structures  $\mathcal{G}$ . Algorithm 1 and 2 presents an instantiation of the proposed SONew method, which solves (6) using  $\mathcal{O}(n)$  time and memory for banded matrices with band size  $b$ . In particular a tridiagonal matrix, corresponding to a chain graph, is a banded matrix with bandsize 1.
Algorithm 1 Sparsified Online Newton (SONew) Algorithm
Inputs:  $\lambda_{t}\coloneqq$  coefficient in the update (10),  $\mathcal{G}\coloneqq$  sparsity graph (banded/tridiagonal),  $\epsilon \coloneqq$  damping parameter,  $T\coloneqq$  total number of iterations/mini-batches,  $\eta_t\coloneqq$  step size/learning rate. Output:  $w_{T + 1}$    
1:  $H_0 = \epsilon I_d,w_1 = 0$    
2: for  $t\in \{1,\dots ,T\}$  do   
3: compute  $g_{t} = \nabla f_{t}(w_{t})$    
4:  $H_{t}\coloneqq H_{t - 1} + P_{\mathcal{G}}(g_{t}g_{t}^{T} / \lambda_{t})\in S_{n}(\mathcal{G})$  with  $P_{\mathcal{G}}$  as in (8).  $\triangleright \mathcal{O}(n)$  time & memory   
5: Get  $L,D =$  SPARSIFIED_INVERSE  $(H_t,\mathcal{G})$  where  $X_{t} = LDL^{T}$  solves (11).   
6: Compute descent direction  $u_{t} = LDL^{T}g_{t}$    
7:  $w_{t + 1} = w_t - \eta_t u_t$    
8: end for   
9: return  $w_{T + 1}$
Algorithm 2 SPARSIFIED_INVERSE  $(H,\mathcal{G})$  in  $\mathcal{O}(n)$  flops
Inputs:  $H\in S_n(\mathcal{G})$  , is as (10).  $\mathcal{G}\coloneqq$  the banded graph of band size  $b\ll n$  Outputs: lower triangular banded  $L\in \mathbb{R}^{n\times n}$  and diagonal matrix  $D\in \mathbb{R}^{n\times n}$    
1: function SPARSIFIED_INVERSE(H, G)   
2:  $L\coloneqq 0,D\coloneqq 0$    
3:  $L_{jj}\coloneqq 1,\forall j\in [n]$    
4: for  $j\in \{1,\dots ,n\}$  do ▷ parallelizable   
5: Let  $H_{jI_j}$  and  $H_{I_jI_j}$  be defined as in Section 2, where  $I_{j}$ $= \{j + 1,\ldots ,j + b\} \cap [n],$    
6: Solve for  $L_{I_{j}j}$  in the linear system  $H_{I_jI_j}L_{I_jj} = -H_{I_jj}\triangleright \mathcal{O}(b^3)$  time.   
7:  $D_{jj}\coloneqq 1 / (H_{jj} + H_{I_{jj}}^T L_{I_{jj}})$    
8: end for   
9: return  $L,D$    
10: end function
Maintaining  $H_{t} \in S_{n}(\mathcal{G})$  in line 4. Solving the subproblem in (6) naively is impractical since  $X_{t-1}^{-1}$  is a dense matrix. However, the structure of the LogDet divergence comes to the rescue; the optimization problem in (6) can be expanded as follows:
$$
\underset {X \in S _ {n} (\mathcal {G}) ^ {+ +}} {\arg \min } - \log \det  (X) + \operatorname {T r} \left(X \left(X _ {t - 1} ^ {- 1} + g _ {t} g _ {t} ^ {T} / \lambda_ {t}\right)\right). \tag {7}
$$
Let us define the projection onto  $S_{n}(\mathcal{G})$ ,  $P_{\mathcal{G}}:\mathbb{R}^{n\times n}\to \mathbb{R}^{n\times n}$  as:
$$
P _ {\mathcal {G}} (M) _ {i j} = \left\{ \begin{array}{l l} M _ {i j} & \text {i f} (i, j) \in E _ {\mathcal {G}}, \\ 0 & \text {o t h e r w i s e .} \end{array} \right. \tag {8}
$$
Note that the  $\mathrm{Tr}(.)$  term in (7) is dependent only on the non-zero elements of  $X\in S_n(\mathcal{G})^{++}$ , since  $\mathrm{Tr}(AB) = \langle A,B\rangle$ , for symmetric matrices  $A$  and  $B$ . Hence, (7) can be written as
$$
\underset {X \in S _ {n} (\mathcal {G}) ^ {+ +}} {\arg \min } - \log \det  (X) + \langle X, P _ {\mathcal {G}} \left(X _ {t - 1} ^ {- 1} + g _ {t} g _ {t} ^ {T} / \lambda_ {t}\right) \rangle , \tag {9}
$$
Computing the entire matrix  $X_{t - 1}^{-1}$  can be avoided by analyzing the optimality condition of (9). Let  $g(X) = -\log \det(X) + \langle X, P_{\mathcal{G}}(X_{t - 1}^{-1} + g_t g_t^T / \lambda_t) \rangle$  denote the objective function in (9), then the optimality condition of (9) is  $P_{\mathcal{G}}(\nabla g(X)) = P_{\mathcal{G}}(\nabla (-\log \det(X) + \langle X, P_{\mathcal{G}}(X_{t - 1}^{-1} + g_t g_t^T / \lambda_t) \rangle) = 0,$  since gradients with respective nonzero entries of  $X$  should be zero,  $\frac{\partial g(X)}{\partial X_{i,j}} = (\nabla_X(g(X)))_{i,j} = 0,\forall (i,j) \in E_{\mathcal{G}}$ . Using  $\nabla(-\log \det(X)) = -X^{-1}, \nabla_X(\langle X,Y \rangle) = Y$ , and setting  $X = X_t$  gives:
$$
P _ {\mathcal {G}} (X _ {t} ^ {- 1}) - P _ {\mathcal {G}} (X _ {t - 1} ^ {- 1} + g _ {t} g _ {t} ^ {T} / \lambda_ {t}) = 0,
$$
$$
H _ {t} = H _ {t - 1} + P _ {\mathcal {G}} \left(g _ {t} g _ {t} ^ {T} / \lambda_ {t}\right), \quad \text {w h e r e} H _ {t} = P _ {\mathcal {G}} \left(X _ {t} ^ {- 1}\right) \tag {10}
$$
Thus we only need to maintain  $H_{t} = P_{\mathcal{G}}(X_{t}^{-1})$ . This matrix is updated as  $H_{t} = H_{t - 1} + P_{\mathcal{G}}(g_{t}g_{t}^{T} / \lambda_{t})$ . Since  $H_{t}\in S_{n}(\mathcal{G})$ , the update can be done in  $\mathcal{O}(|E_{\mathcal{G}}|)$  memory and time, while computing the matrix  $X_{t}^{-1}$  would have cost  $\mathcal{O}(n^2)$ . In SONew (Algorithm 1), this key observation is used to maintain  $H_{t}$  in line 4.
Computing  $X_{t}$  in line 5. Now that  $H_{t}$  is known at every round  $t$ , we can replace  $P_{\mathcal{G}}(X_{t - 1}^{-1} + g_t g_t^T / \lambda_t)$  in (9) with  $H_{t}$  as:
$$
X _ {t} = \underset {X \in S _ {n} (\mathcal {G}) ^ {+ +}} {\arg \min } - \log \det  (X) + \operatorname {T r} \left(X H _ {t}\right). \tag {11}
$$
For an arbitrary graph  $\mathcal{G}$ , solving (11) might be difficult. Theorems 3.1 and 3.2 show embarrassingly parallelizable explicit solutions to the subproblem (11) for tridiagonal and banded sparsity patterns.
Theorem 3.1 (Explicit solution of (11) for tridiagonal structures/chain graph). Let the sparsity structure  $\mathcal{G}$  be a chain with edges  $E_{\mathcal{G}} = \{(i,j):|i - j|\leq 1,1\leq i,j\leq n\}$ . Also, let  $H\in S_n(\mathcal{G})$  be such that any submatrix of  $H$  corresponding to a complete subgraph of  $\mathcal{G}$  is positive definite, then the solution of (11) is given by  $\hat{X} = LDL^{T}$ , where the unit lower bidiagonal matrix  $L$  and diagonal matrix  $D$  have the following non-zero entries: