Title: LOW-RANK TRAINING OF DEEP NEURAL NETWORKS FOR EMERGING MEMORY TECHNOLOGY Anonymous

Abstract: The recent success of neural networks for solving difficult decision tasks has incentivized incorporating smart decision making "at the edge." However, this work has traditionally focused on neural network inference, rather than training, due to memory and compute limitations, especially in emerging non-volatile memory systems, where writes are energetically costly and reduce lifespan. Yet, the ability to train at the edge is becoming increasingly important as it enables real-time adaptability to device drift and environmental variation, user customization, and federated learning across devices. In this work, we address two key challenges for training on edge devices with non-volatile memory: low write density and low auxiliary memory. We present a low-rank training scheme that addresses these challenges while maintaining computational efficiency. We then demonstrate the technique on a representative convolutional neural network across several adaptation problems, where it out-performs standard SGD both in accuracy and in number of weight writes.

1. INTRODUCTION:

Deep neural networks have shown remarkable performance on a variety of challenging inference tasks. As the energy efficiency of deep-learning inference accelerators improves, some models are now being deployed directly to edge devices to take advantage of increased privacy, reduced network bandwidth, and lower inference latency. Despite edge deployment, training happens predominately in the cloud. This limits the privacy advantages of running models on-device and results in static models that do not adapt to evolving data distributions in the field.

Efforts aimed at on-device training address some of these challenges. Federated learning aims to keep data on-device by training models in a distributed fashion (Konecný et al., 2016) . On-device model customization has been achieved by techniques such as weight-imprinting (Qi et al., 2018) , or by retraining limited sets of layers. On-chip training has also been demonstrated for handling hardware imperfections (Zhang et al., 2017; Gonugondla et al., 2018) . Despite this progress with small models, on-chip training of larger models is bottlenecked by the limited memory size and compute horsepower of edge processors.

Emerging non-volatile (NVM) memories such as resistive random access memory (RRAM) have shown great promise for energy and area-efficient inference (Yu, 2018) . However, on-chip training requires a large number of writes to the memory, and RRAM writes cost significantly more energy than reads (e.g., 10.9 pJ/bit versus 1.76 pJ/bit (Wu et al., 2019) ). Additionally, RRAM endurance is on the order of 10 6 writes (Grossi et al., 2019) , shortening the lifetime of a device due to memory writes for on-chip training.

In this paper, we present an online training scheme amenable to NVM memories to enable next generation edge devices. Our contributions are (1) an algorithm called Streaming Kronecker Sum Approximation (SKS), and its analysis, which addresses the two key challenges of low write density and low auxiliary memory; (2) two techniques "gradient max-norm" and "streaming batch norm" to help training specifically in the online setting; (3) a suite of adaptation experiments to demonstrate the advantages of our approach.

2. RELATED WORK:

Efficient training for resistive arrays. Several works have aimed at improving the efficiency of training algorithms on resistive arrays. Of the three weight-computations required in training (forward, backprop, and weight update), weight updates are the hardest to parallelize using the array structure. Stochastic weight updates (Gokmen & Vlasov, 2016) allow programming of all cells in a crossbar at once, as opposed to row/column-wise updating. Online Manhattan rule updating (Zamanidoost et al., 2015) can also be used to update all the weights at once. Several works have proposed new memory structures to improve the efficiency of training (Soudry et al., 2015; Ambrogio et al., 2018) . The number of writes has also been quantified in the context of chip-in-the-loop training (Yu et al., 2016) .

Distributed gradient descent. Distributed training in the data center is another problem that suffers from expensive weight updates. Here, the model is replicated onto many compute nodes and in each training iteration, the mini-batch is split across the nodes to compute gradients. The distributed gradients are then accumulated on a central node that computes the updated weights and broadcasts them. These systems can be limited by communication bandwidth, and compressed gradient techniques (Aji & Heafield, 2017) have therefore been developed. In Lin et al. (2017) , the gradients are accumulated over multiple training iterations on each compute node and only gradients that exceed a threshold are communicated back to the central node. In the context of on-chip training with NVM, this method helps reduce the number of weight updates. However, the gradient accumulator requires as much memory as the weights themselves, which negates the density benefits of NVM.

Low-Rank Training. Our work draws heavily from previous low-rank training schemes that have largely been developed for use in recurrent neural networks to uncouple the training memory requirements from the number of time steps inherent to the standard truncated backpropagation through time (TBPTT) training algorithm. Algorithms developed since then to address the memory problem include Real-Time Recurrent Learning (RTRL) (Williams & Zipser, 1989) , Unbiased Online Recurrent Optimization (UORO) (Tallec & Ollivier, 2017) , Kronecker Factored RTRL (KF-RTRL) (Mujika et al., 2018) , and Optimal Kronecker Sums (OK) (Benzing et al., 2019) . These latter few techniques rely on the weight gradients in a weight-vector product looking like a sum of outer products (i.e., Kronecker sums) of input vectors with backpropagated errors. Instead of storing a growing number of these sums, they can be approximated with a low-rank representation involving fewer sums.

3. TRAINING NON-VOLATILE MEMORY:

The meat of most deep learning systems are many weight matrix -activation vector products W • a. Fully-connected (dense) layers use them explicitly:

a [ ] = σ W [ ] a [ -1] + b [ ]

for layer , where σ is a non-linear activation function (more details are discussed in detail in Appendix C.1). Recurrent neural networks use one or many matrix-vector products per recurrent cell. Convolutional layers can also be interpreted in terms of matrix-vector products by unrolling the input feature map into strided convolution-kernel-size slices. Then, each matrix-vector product takes one such input slice and maps it to all channels of the corresponding output pixel (more details are discussed in Appendix C.2).

The ubiquity of matrix-vector products allows us to adapt the techniques discussed in "Low-Rank Training" of Section 2 to other network architectures. Instead of reducing the memory across time steps, we can reduce the memory across training samples in the case of a traditional feedforward neural network. However, in traditional training (e.g., on a GPU), this technique does not confer advantages. Traditional training platforms often have ample memory to store a batch of activations and backpropagated gradients, and the weight updates ∆W can be applied directly to the weights W once they are computed, allowing temporary activation memory to be deleted. The benefits of low-rank training only become apparent when looking at the challenges of proposed NVM devices:

Low write density (LWD). In NVM, writing to weights at every sample is costly in energy, time, and endurance. These concerns are exacerbated in multilevel cells, which require several steps of an iterative write-verify cycle to program the desired level. We therefore want to minimize the number of writes to NVM.

Low auxiliary memory (LAM):

. NVM is the densest form of memory. In 40nm technology, RRAM 1T-1R bitcells @ 0.085 um 2 (Chou et al., 2018) are 2.8x smaller than 6T SRAM cells @ 0.242 um2 (TSMC, 2019). Therefore, NVM should be used to store the memory-intensive weights. By the same token, no other on-chip memory should come close to the size of the on-chip NVM. In particular, if our b-bit NVM stores a weight matrix of size n o × n i , we should use at most r(n i + n o )b auxiliary non-NVM memory, where r is a small constant. Despite these space limitations, the reason we might opt to use auxiliary (large, high endurance, low energy) memory is because there are places where writes are frequent, violating LWD if we were to use NVM.

In the traditional minibatch SGD setting with batch size B, an upper limit on the write density per cell per sample is easily seen: 1/B. However, to store such a batch of updates without intermediate writes to NVM would require auxiliary memory proportional to B. Therefore, a trade-off becomes apparent. If B is reduced, LAM is satisfied at the cost of LWD. If B is raised, LWD is satisfied at the cost of LAM. Using low-rank training techniques, the auxiliary memory requirements are decoupled from the batch size, allowing us to increase B while satisfying both LWD and LAM1 . Additionally, because the low-rank representation uses so little memory, a larger bitwidth can be used, potentially allowing for gradient accumulation in a way that is not possible with low bitwidth NVM weights. In the next section, we elaborate on the low-rank training method.

4. LOW-RANK TRAINING METHOD:

Let z (i) = W a (i) + b be the standard affine transformation building block of some larger network, e.g., y i) . A minibatch SGD weight update accumulates this gradient over B samples:

(i) p = f post (z (i) ) and a (i) = f pre (x (i) ) with prediction loss L(y (i) p , y (i) t ), where (x (i) , y (i) t ) is the i th training sample pair. Then weight gradient ∇ W L (i) = dz (i) a (i) = dz (i) ⊗ a (i) where dz (i) = ∇ z (i) L (

∆W = -η B i=1 dz (i) ⊗ a (i) for learning rate η.

For a rank-r training scheme, approximate the sum

B i=1 dz (i) ⊗a (i) by iteratively updating two rank- r matrices L ∈ R no×r , R ∈ R ni×r with each new outer product: L R ← rankReduce( L R + dz (i) ⊗ a (i) )

. Therefore, at each sample, we convert the rank-q = r + 1 system L R + dz (i) ⊗ a (i) into the rank-r L R . In the next sections, we discuss how to compute rankReduce.

4.1. OPTIMAL KRONECKER SUM APPROXIMATION (OK):

One option for rankReduce(X) to convert from rank q = r + 1 X to rank r is a minimum error estimator, which is implemented by selecting the top r components of a singular value decomposition (SVD) of X. However, a naïve implementation is computationally infeasible and biased: E[rankReduce(X)] = X. Benzing et al. (2019) solves these problems by proposing a minimum variance unbiased estimator for rankReduce, which they call the OK algorithm 2 .

The OK algorithm can be understood in two key steps: first, an efficient method of computing the SVD of a Kronecker sum; second, a method of splitting the singular value matrix Σ into two rank-r matrices whose outer product is a minimum-variance, unbiased estimate of Σ. Details can be found in their paper, however we include a high-level explanation in Sections 4.1.1 and 4.1.2 to aid our discussions. Note that our variable notation differs from Benzing et al. (2019) .

4.1.1. EFFICIENT SVD OF KRONECKER SUMS:

Let L = [ L, dz (i) ] and R = [ R, a (i) ] so that LR = L R +dz (i) ⊗a (i) . Recall that rankReduce should turn rank-q LR into an updated rank-r L R . & Dhillon, 2006) , making it computationally feasible on small devices. Now we have:

QR-factorize:

L = Q L R L and R = Q R R R where Q L ∈ R no×q , Q R ∈ R ni×q are orthogonal so that LR = Q L (R L R R )Q R . Let C = R L R R ∈ R q×q . Then we can find the SVD of C = U C ΣV C in O(q 3 ) time (Cline

LR = Q L (U C ΣV C )Q R = (Q L U C )Σ(Q R V C ) (1)

which gives the SVD of LR since Q L U C and Q R V C are orthogonal and Σ is diagonal. This SVD computation has a time complexity of O((n i +n o +q)q 2 ) and a space complexity of O((n i +n o +q)q).

4.1.2. MINIMUM VARIANCE, UNBIASED ESTIMATE OF Σ:

In Benzing et al. (2019) , it is shown that the problem of finding a rank-r minimum variance unbiased estimator of LR can be reduced to the problem of finding a rank-r minimum variance unbiased estimator of Σ and plugging it in to (1).

Further, it is shown that such an optimal approximator for Σ = diag(σ 1 , σ 2 , . . . , σ q ), where σ 1 ≥ σ 2 ≥ • • • ≥ σ q will involve keeping the m -1 largest singular values and mixing the smaller singular values σ m , . . . , σ q within their (k + 1) × (k + 1) submatrix with m, k defined below. Let: k) be orthogonal such that its left nullspace is the span of x 0 . Then XX = I -x 0 x 0 . Now, let s ∈ {-1, 1} (k+1)×1 be uniform random signs and define:

m = min i s.t. (q -i)σ i ≤ q j=i σ j k = q -m x 0 = 1 - σ m k s 1 , . . . , 1 - σ q k s 1 s 1 = q i=m σ i Note that ||x 0 || 2 = 1. Let X ∈ R (k+1)×(

EQUATION

where is an element-wise product. Then ΣL Σ R = Σ is a minimum variance, unbiased3 rank-r approximation of Σ. Plugging Σ into (1),

LR = (Q L U C )Σ(Q R V C ) ≈ (Q L U C ) Σ(Q R V C ) = (Q L U C ΣL )(Q R V C ΣR ) (3) Thus, L = Q L U C ΣL ∈ R no×r and R = Q R V C ΣR ∈ R ni×r

gives us a minimum variance, unbiased, rank-r approximation L R .

4.2. STREAMING KRONECKER SUM APPROXIMATION (SKS):

Although the standalone OK algorithm presented by Benzing et al. (2019) has good asymptotic computational complexity, our vector-vector outer product sum use case permits further optimizations.

In this section we present these optimizations, and we refer readers to the explicit implementation called Streaming Kronecker Sum Approximation (SKS) in Algorithm 1 of Appendix A.

4.2.1. MAINTAIN ORTHOGONAL:

Q L , Q R

The main optimization is a method of avoiding recomputing the QR factorization of L and R at every step. Instead, we keep track of orthogonal matrices Q L , Q R , and weightings c x such that

L = Q L • diag( √ c x ) [:r] and R = Q R • diag( √ c x ) [:r] .

Upon receiving a new sample, a single inner loop of the numerically-stable modified Gram-Schmidt (MGS) algorithm (Björck, 1967) can be used to update Q L and Q R . The orthogonal basis coefficients c L = Q L dz (i) and c R = Q R a (i) computed during MGS can be used to find the new value of C = c L c R + diag(c x ).

After computing ΣL = ΣR in (2), we can orthogonalize these matrices into

ΣL = ΣR = Q x R x . Then from (3), we have L R = (Q L U C Q x )(R x R x )(Q R V C Q x )

. With this formulation, we can maintain orthogonality in Q L , Q R by setting:

Q L ← Q L U C Q x Q R ← Q R V C Q x c x ← diag(R x R x )

These matrix multiplies require O((n i +n o )q 2 ) multiplications, so this optimization does not improve asymptotic complexity bounds. This optimization may nonetheless be practically significant since matrix multiplies are easy to parallelize and would typically not be the bottleneck of the computation compared to Gram-Schmidt. The next section discusses how to orthogonalize ΣL efficiently and why (R x R x ) is diagonal.

4.2.2. ORTHOGONALIZATION OF ΣL:

Orthogonalization of ΣL is relatively straightforward. From (2), the columns of ΣL are orthogonal since Z is orthogonal. However, they do not have unit norm. We can therefore pull out the norm into a separate diagonal matrix R x with diagonal elements √ c x :

Q x = I m-1 0 0 X s √ c x = ( √ σ 1 , . . . , √ σ m-1 , s 1 /k q-m+1 times ) 4.2.3 FINDING ORTHONORMAL BASIS X

We generated X by finding an orthonormal basis that was orthogonal to a vector x 0 so that we could have XX = I -x 0 x 0 . An efficient method of producing this basis is through Householder matrices Householder, 1958; user1551, 2013) .

(x 0 , X) = I -2 vv /||v|| 2 where v = x 0 -e (1) and (x 0 , X) is a k + 1 × k + 1 matrix with first column x 0 and remaining columns X (

4.2.4. EFFICIENCY COMPARISONS TO STANDARD APPROACH:

The OK/SKS methods require O((n i + n o + q)q 2 ) operations per sample and O(n i n o q) operations after collecting B samples, giving an amortized cost of O((n i + n o + q)q 2 + n i n o q/B) operations per sample. Meanwhile, a standard approach expands the Kronecker sum at each sample, costing O(n i n o ) operations per sample. If q B, n i , n o then the low rank method is superior to minibatch SGD in both memory and computational cost.

5. CONVEX CONVERGENCE:

SKS introduces variance into the gradient estimates, so here we analyze the implications for online convex convergence. We analyze the case of strongly convex loss landscapes f t (w t ) for flattened weight vector w t and online sample t. In Appendix B, we show that with inverse squareroot learning rate, when the loss landscape Hessians satisfy 0 ≺ cI ∇ 2 f t (w t ) and under constraint (4) for the size of gradient errors ε t , where w * is the optimal offline weight vector, the online regret ( 5) is sublinear in the number of online steps T . We can approximate ||ε|| and show that convex convergence is likely when ( 6) is satisfied in the biased, zero-variance case (equivalent to raw SVD, i.e., not applying Section 4.1.2), or when ( 7) is satisfied in the unbiased, minimum-variance case.

||ε t || ≤ c 2 ||w t -w * || (4) R(T ) = T t=1 f t (w t ) - T t=1 f t (w * ) (5) B i=1 σ (t,i) q 2 ≤ c 2 4 ||w t -w * || 2 (6) B i=1 σ (t,i) r σ (t,i) q ≤ c 2 8 ||w t -w * || 2 (7)

Equations (6, 7) suggest conditions under which fast convergence may be more or less likely and also point to methods for improving convergence. We discuss these in more detail in Appendix B.3.

5.1. CONVERGENCE EXPERIMENTS:

We validate (4) with several linear regression experiments on a static input batch X ∈ R 1024×100 and target Y t ∈ R 256×100 . In Figure 1 (a), Gaussian noise at different strengths (represented by different colors) is added to the true batch gradients at each update step. Notice that convergence slows significantly to the right of the dashed lines, which is the region where (4) no longer holds4 .

In Figure 1 (b), we validate Equations (4, 6, 7) by testing the SVD and SKS cases with rank r = 10. In these particular experiments, SKS adds too much variance, causing it to operate to the right of the dashed lines. However, both SVD and SKS can be seen to reduce their variance as training progresses.

In the case of SVD, it is able to continue training as it tracks the right dashed line.

10 2 10 0 10 2 10 4 106 10 8 10 10 10 12

Gradient Error Variance 

6. IMPLEMENTATION DETAILS:

Quantization. The NN is quantized in both the forward and backward directions with uniform power-of-2 quantization, where the clipping ranges are fixed at the start of training5 . Weights are quantized to 8 bits between -1 and 1, biases to 16 bits between -8 and 8, activations to 8 bits between 0 and 2, and gradients to 8 bits between -1 and 1. Both the weights W and weight updates ∆W are quantized to the same LSB so that weights cannot be used for accumulation beyond the fixed quantization dynamic range. This is in contrast to using high bitwidth (Zhou et al., 2016; Banner et al., 2018) or floating point accumulators. See Appendix D for more details on quantization.

Gradient Max-Norming. State-of-the-art methods in training, such as Adam (Kingma & Ba, 2014), use auxiliary memory per parameter to normalize the gradients. Unfortunately, we lack the memory budget to support these additional variables, especially if they must be updated every sample 6 . Instead, we propose dividing each gradient tensor by the maximum absolute value of its elements. This stabilizes the range of gradients across samples. See Appendix E for more details on gradient max-norming. In the experiments, we refer to this method as "max-norm" (opposite "no-norm").

Streaming Batch Normalization. Batch normalization (Ioffe & Szegedy, 2015) is a powerful technique for improving training performance which has been suggested to work by smoothing the loss landscape (Santurkar et al., 2018) . We hypothesize that this may be especially helpful when parameters are quantized as in our case. However, in the online setting, we receive samples one-at-a-time rather than in batches. We therefore propose a streaming batch norm that uses moving average statistics rather than batch statistics as described in detail in Appendix F.

7. EXPERIMENTS 7.1 ADAPTATION EXPERIMENTS:

To test the effectiveness of SKS, experiments are performed on a representative CNN with four 3 × 3 convolution layers and two fully-connected layers. We generate "offline" and "online" datasets based on MNIST (see Appendix G), including one in which the statistical distribution shifts every 10k images. We then optimize an online SGD and rank-4 SKS model for fair comparison (see Appendix H). To see the importance of different training techniques, we run several ablations in Appendix I. Finally, we compare these different training schemes in different environments, meant to model real life. In these hypothetical scenarios, a model is first trained on the offline training set, and is then deployed to a number of devices at the edge that make supervised predictions (they make a prediction, then are told what the correct prediction would have been).

We present results on four hypothetical scenarios. 

SGD:

-+0.3 ± 0.2 +0.3 ± 0.2 +0.3 ± 0.2 +0.9 ± 0.2 -3.9 ± 0.8 UORO 1 +0.4 ± 0.2 +0.3 ± 0.4 -1.8 ± 0.9 -7.6 ± 1.6 -31.7 ± 1.6 SVD 1 +1.9 ± 0.2 +5.8 ± 1.0 -3.4 ± 1.0 -19.4 ± 0.9 -40.7 ± 1.1 2 +1.4 ± 0.4 +6.5 ± 0.7 +6.3 ± 0.6 -5.2 ± 0.9 -36.3 ± 0.9 4 +1.3 ± 0.4 +6.5 ± 0.7 +5.2 ± 0.8 -3.3 ± 1.0 -33.8 ± 0.8 8 +1.4 ± 0.3 +5.6 ± 0.8 +4.3 ± 0.9 -2.4 ± 1.0 -32.8 ± 0.9

SKS 1 +0.3 ± 0.2 +0.3 ± 0.2 -0.7 ± 0.4 -2.7 ± 1.7 -26.5 ± 2.6 2 +0.3 ± 0.2 +0.4 ± 0.3 -0.1 ± 0.4 +1.3 ± 0.9 -12.9 ± 1.1 4 +0.4 ± 0.2 +0.6 ± 0.2 +1.9 ± 0.3 +8.0 ± 1.1 -5.1 ± 1.1 8 +0.4 ± 0.2 +1.1 ± 0.2 +3.3 ± 0.7 +4.8 ± 1.5 -15.8 ± 1.7

8. CONCLUSION:

We demonstrated the potential for SKS to solve the major challenges facing online training on NVM-based edge devices: low write density and low auxiliary memory. SKS is a computationallyefficient, memory-light algorithm capable of decoupling batch size from auxiliary memory, allowing larger effective batch sizes, and consequently lower write densities. Additionally, we noted that SKS may allow for training under severe weight quantization constraints as rudimentary gradient accumulations are handled by the L, R matrices, which can have high bitwidths (as opposed to SGD, which may squash small gradients to 0). We found expressions for when SKS might have better convergence properties. Across a variety of online adaptation problems and a large-scale transfer learning demonstration, SKS was shown to match or exceed the performance of SGD while using a small fraction of the number of updates. Finally, we suspect that these techniques could be applied to a broader range of problems. Auxiliary memory minimization may be analogous to communication minimization in training strategies such as federated learning, where gradient compression is important.

A SKS ALGORITHM:

Algorithm 1 Streaming Kronecker Sum Approximation

State: Q L ∈ R no×q ; Q R ∈ R ni×q ; c x ∈ R q×1 Input: dz (i) ∈ R no×1 ; a (i) ∈ R ni×1 for i ∈ [1, B] for i = 1 . . . B do {Modified Gram-Schmidt.} c L , c R ← 0 q×1 for j = 1 . . . r do c L,j ← Q L,j • dz (i) ; dz (i) ← dz (i) -c L,j • Q L,j c R,j ← Q R,j • a (i) ; a (i) ← a (i) -c L,j • Q L,j end for c L,q ← ||dz (i) ||; Q L,q ← dz (i) /c L,q c R,q ← ||a (i) ||; Q R,q ← a (i) /c R,q {Generate C and find its SVD.} C ← c L c R + diag(c x ) U C • diag(σ) • V C ← SVD(C) {Minimum-variance unbiased estimator for Σ.} m ← min j s.t. (q -j)σ j ≤ q =j σ s 1 ← q i=m σ i , k ← q -m v ← 1 -k/s 1 • σ [m:] -e (1) s ← {-1, 1} (k+1)×1 {Ind. uniform random signs.} X s ← I + (s v)(v/v 1 ) [2:] {Householder.} {QR-factorization of ΣL .} Q x ← I 0 0 X s ∈ R q×r c x ← (σ 1 , . . . , σ m-1 , s 1 /k, . . . , s 1 /k q-m+1 times ) {Update the first r columns of Q L , Q R .} Q L[:r] ← Q L • U C • Q x Q R[:r] ← Q R • V C • Q x end for {Compute final L, R where ∇ W L ≈ L R .} L ← Q L • diag( √ c x ) [:r] R ← Q R • diag( √ c x ) [:r]

B CONVEX CONVERGENCE:

In this section we will attempt to bound the regret (defined below) of an SGD algorithm using noisy SKS estimates g = g + ε in the convex setting, where g are the true gradients and ε are the errors introduced by the low rank SKS approximation. Here, g is a vector of size N and can be thought of as a flattened/concatenated version of the gradient tensors (e.g., N = n i • n o ).

Our proof follows the proof in Zinkevich (2003) . We define F as the convex feasible set (valid settings for our weight tensors) and assume that F is bounded with D = max w,v∈F ||w -v|| being the maximum distance between two elements of F. Further, assume a batch t of B samples out of T total batches corresponds to a loss landscape f t (w t ) that is strongly convex in weight parameters w t , so there are positive constants C ≥ c > 0 such that cI ∇ 2 f t (w t ) CI for all t (Boyd & Vandenberghe, 2004, Section 9.3). We define regret as R(T ) = T t=1 f t (w t ) -T t=1 f t (w * ) where w * = argmin w T t=1 f t (w) (i.e., it is an optimal offline minimizer of

f 1 , • • • , f T ).

The gradients seen during SGD are g t = ∇f t (w t ) and we assume they are bounded by G = max w∈F,t∈[1,T ] ||∇f t (w)||. We also assume errors are bounded by

E = max t∈[1,T ] ||ε t ||. Therefore, max t∈[1,T ] ||g t || ≤ max t∈[1,T ] ||g t || + ||ε t || ≤ G + E by the triangle inequality.

Theorem 1. Assume SKS-based SGD is applied with learning rate η t = 1/ √ t. Then, under the additional constraint g t • (w tw * ) -c 2 ||w tw * || 2 2 ≤ gt • (w tw * ), we have sublinear regret:

R(T ) ≤ D 2 2 √ T + (G + E) 2 √ T -1/2

Proof. From strong convexity cI ∇ 2 f t (w t ) CI for all t,

EQUATION

In particular, if we consider v = w * and rearrange,

EQUATION

Consider a gradient update w t+1 = P F (w t -η t gt ), where P F projects the update back to F. Then,

EQUATION

From (9, 10),

EQUATION

We now bound the regret:

R(T ) = T t=1 f t (w t ) -f t (w * ) ≤ T t=1 1 2η t ||w t -w * || 2 2 -||w t+1 -w * || 2 2 + η t 2 (G + E) 2 = ||w 1 -w * || 2 2 2η 1 - ||w T +1 -w * || 2 2 2η T + 1 2 T t=2 1 η t - 1 η t-1 ||w t -w * || 2 2 + (G + E) 2 2 T t=1 η t ≤ ||w 1 -w * || 2 2 2η 1 + 1 2 T t=2 1 η t - 1 η t-1 ||w t -w * || 2 2 + (G + E) 2 2 T t=1 η t ≤ D 2 2η 1 + 1 2 T t=2 1 η t - 1 η t-1 D 2 + (G + E) 2 2 T t=1 η t = D 2 2η T + (G + E) 2 2 T t=1 η t (12)

Under review as a conference paper at ICLR 2020 Zinkevich, 2003) , so from (12),

If η t = 1/ √ t, then T t=1 η t ≤ 2 √ T -1 (

R(T ) ≤ D 2 2 √ T + (G + E) 2 √ T -1/2 (13)

This is a sublinear regret and therefore, average regret R(T )/T is bounded above by 0 in the limit as T → ∞. To achieve this result, we constrained g t • (w tw * ) -c 2 ||w tw * || 2 2 ≤ gt • (w tw * ). We now examine sufficient conditions for this inequality to be satisfied.

g t • (w t -w * ) - c 2 ||w t -w * || 2 2 ≤ gt • (w t -w * ) ε t • (w * -w t ) ≤ c 2 ||w t -w * || 2 2 (14) Since ε t • (w * -w t ) ≤ ||ε t || • ||w * -w t

|| by Cauchy-Schwarz, it is sufficient for:

EQUATION

B.1 CONSIDERATIONS FOR RANK DEFICIENT HESSIANS:

In the preceding proof, we assumed c > 0. However, it is common for this to not hold. For example, in linear regression, where c = λ min (XX ) for sample input X ∈ R (ni×B) (Haykin, 2014, Chapter 4 .3), if B < n i then c = 0. We can modify (15) to handle this case. Let c be the minimum non-zero Eigenvalue of XX and let w ∈ R B represent w ∈ R N in the Eigenbasis of XX . Then (15) becomes:

EQUATION

B.2 ESTIMATES FOR THE SKS ERROR:

We can estimate the SKS error ||ε t || in both the biased and unbiased cases. For the biased, zerovariance case, we get rid of the lowest singular value (out of q singular values) as we see each sample. Thus, at a given sample i, the error is σ

(i) q • Q (i) L,q • Q (i)

R,q and the average squared error is (1/N )σ

(i)2

q . We can treat this as a per-element variance. If these smallest singular components are uncorrelated from sample to sample, then the variances add:

σ 2 ε ≈ 1 N B i=1 σ (i) q 2 (17)

For the unbiased, minimum-variance case, Theorem A.4 from Benzing et al. (2019) states that the minimum variance is s 2 1 /k + s 2 where s 1 = q i=m σ i , s 2 = q i=m σ 2 i , and k, m are as defined in Section 4.1.2. Since m is chosen to minimize variance, we can upper bound the variance by choosing m = r and therefore k = 1, s 1 = σ r + σ q , and s 2 = σ 2 r + σ 2 q . Empirically, this tends to be a good approximation. Then, the average per-element variance added at sample i is approximately

(2/N ) σ (i) r σ (i) q

. Assuming errors between samples are uncorrelated, this leads to a total variance:

σ 2 ε ≈ 2 N B i=1 σ (i) r σ (i) q (18)

For either case, ||ε|| 2 ≈ N σ 2 ε . For the t-th batch and i-th sample, we denote σ (t,i) q

as the q-th singular value. For simplicity, we focus on the biased, zero-variance case (the unbiased case is similar). From ( 15), an approximately sufficient condition for sublinear-regret convergence is:

B i=1 σ (t,i) q 2 ≤ c 2 4 ||w t -w * || 2 (19) B.3 DISCUSSION ON CONVERGENCE

Equation ( 19) suggests that as w t → w * , the constraints for achieving sublinear-regret convergence become more difficult to maintain. However, in practice this may be highly problem-dependent as the σ q will also tend to decrease near optimal solutions. To get a better sense of the behavior of the left-hand side of ( 19), suppose that: no×ni) are the matrix weight W t gradients at batch t and || • || F is a Frobenius norm. We therefore expect both the left (proportional to ||G t || 2 F ) and the right (proportional to ||w tw * || 2 ) of ( 19) to decrease during training as w t → w * . This behavior is in fact what is seen in Figure 1 (b). If achieving convergence is found to be difficult, ( 19) provides some insight for convergence improvement methods.

B i=1 σ (t,i) q 2 ≈ B i=q σ i (G t ) 2 ≤ B i=1 σ i (G t ) 2 = ||G t || 2 F where G t = ∇ W t f t (W t ) ∈ R (

One solution is to reduce batch size B to satisfy the inequality as necessary. This minimizes the weight updates during more repetitive parts of training while allowing dense weight updates (possibly approaching standard SGD with small batch sizes) during more challenging parts of training.

Another solution is to reduce σ q . One way to do this is to increase the rank r so that the spectral energy of the updates are spread across more singular components. There may be alternate approaches based on conditioning the inputs to shape the distribution of singular values in a beneficial way.

A third method is to focus on c, the lower bound on curvature of the convex loss functions. Perhaps a technique such as weight regularization can increase c by adding constant curvature in all Eigendirections of the loss function Hessian (although this may also increase the LHS of ( 19)). Alternatively, perhaps low-curvature Eigen-directions are less important for loss minimization, allowing us to raise the c that we effectively care about. This latter approach requires no particular action on our part, except the recognition that fast convergence may only be guaranteed for high-curvature directions. This is exemplified in Figure 1(b) , where we can see SVD track the curve for C more so than c.

Finally, we note that this analysis focuses solely on the errors introduced by a floating-point version of SKS. Quantization noise can add additional error into the ε t term. We expect this to add a constant offset to the LHS of ( 19). For a weight LSB ∆, quantization noise has variance ∆ 2 /12, so we desire: 

EQUATION

∇ W L = (∇ z L) dz (∇ W z) a = dz ⊗ a (21)

which is exactly the per-sample Kronecker sum update we saw in linear regression. Thus, at every training sample, we can add (dz (i) ⊗ a (i) ) to our low rank estimate with SKS.

C.2 CONVOLUTIONAL LAYER:

A convolutional layer transforms an input feature map A ∈ R hin×win×cin to an intermediate feature map Z = W kern * A + b ∈ R hout×wout×cout through a 2D convolution * with weight kernel W kern ∈ R cout×k h ×kw×cin . Then it computes an output feature map y = σ(z) where σ is a non-linear activation function.

Convolutions can be interpreted as matrix multiplications through the im2col operation which converts the input feature map A into a matrix A col ∈ R (houtwout)×(k h kwcin) where the i th row is a flattened version of the sub-tensor of a which is dotted with W kern to produce the i th pixel of the output feature map (Ren & Xu, 2015) . We can multiply A col by a flattened version of the kernel, W ∈ R cout×(k h hwcin) to perform the W kern * A convolution operation with a matrix multiplication.

Under the matrix multiplication interpretation, weight gradients can be represented as:

EQUATION

which is the same as h out w out Kronecker sum updates. Thus, at every output pixel j of every training sample i, we can add (dZ

(i) col,j ⊗ A (i)

col,j ) to our low rank estimate with SKS. Note that while we already save an impressive factor of B/q in memory when computing gradients for the dense layer, we save a much larger factor of Bh out w out /q in memory when computing gradients for the convolution layers, making the low rank training technique even more crucial here. However, some care must be taken when considering activation memory for convolutions. For compute-constrained edge devices, image dimensions may be small and result in minimal intermediate feature map memory requirements. However, if image dimensions grow substantially, activation memory could dominate compared to weight storage. Clever dataflow strategies may provide a way to reduce intermediate activation storage even when performing backpropagation9 .

D HARDWARE QUANTIZATION MODEL:

In a real device, operations are expected to be performed in fixed point arithmetic. Therefore, all of our training experiments are conducted with quantization in the loop. Our model for quantization is shown in Figure 3 . The green arrows describe the forward computation. Ignoring quantization for a moment, we would have a = ReLU α W * a -1 + b , where * can represent either a convolution or a matrix multiply depending on the layer type and α is the closest power-of-2 to He initialization (He et al., 2015) . For quantization, we rely on four basic quantizers: Qw, Qb, Qa, Qg, which describe weight quantization, bias and intermediate accumulator quantization, activation quantization, and gradient quantization, respectively. All quantizers use fixed clipping ranges as depicted and quantize uniformly within those ranges to the specified bitwidths.

In the backward pass, follow the orange arrows from δ . Backpropagation follows standard backpropagation rules including using the straight-through estimator (Bengio et al., 2013) for quantizer gradients. However, because we want to perform training on edge devices, these gradients must themselves be quantized. The first place this happens is after passing backward through the ReLU derivitive. The other two places are before feeding back into the network parameters W , b , so that W , b cannot be used to accumulate values smaller than their LSB. Finally, instead of deriving ∆W from a backward pass through the * operator, the SKS method is used.

𝑊 ℓ 𝑏 ℓ 𝛼 ℓ 𝑎 ℓ-1 𝑎 ℓ 𝛿 ℓ 𝛿 ℓ-1 SKS × × lr 𝑄𝑤 8𝑏 -1 𝑡𝑜 1 𝑄𝑏 16𝑏 -8 𝑡𝑜 8 𝑄𝑎 8𝑏 0 𝑡𝑜 2 * × 𝑄𝑏 16𝑏 -8 𝑡𝑜 8 + 𝑄𝑏 16𝑏 -8 𝑡𝑜 8 ReLU 𝑄𝑔 8𝑏 -1 𝑡𝑜 1 𝑄𝑤 8𝑏 -1 𝑡𝑜 1 𝑄𝑏 16𝑏 -8 𝑡𝑜 8

Figure 3 : Signal flow graph for a forward and backward quantized convolutional or dense layer.

SKS collects a -1 , dz for many samples before computing the approximate ∆ W . It accumulates information in two low rank matrices L, R which are themselves quantized to 16 bits with clipping ranges determined dynamically by the max absolute value of elements in each matrix. While SKS accumulates for B samples, leading to a factor of B reduction in the rate of updates to W , b is updated at every sample. This is feasible in hardware because b is small enough to be stored in more expensive forms of memory that have superior endurance and write power performance.

Because of the coarse weight LSB size, weight gradients may be consistently quantized to 0, preventing them from accumulating. To combat this, we only apply an update if a minimum update density ρ min = 0.01 would be achieved, otherwise we continue accumulating samples in L and R, which have much higher bitwidths. When an update does finally happen, the "effective batch size" will be a multiple of B and we increase the learning rate correspondingly. In the literature, a linear scaling rule is suggested (see Goyal et al. (2017) ), however we empirically find square-root scaling works better (see Appendix H). Figure 4 plots the magnitude of gradients seen in a weight tensor over training steps. One apparent property of these gradients is that they have a large dynamic range, making them difficult to quantize.

E GRADIENT MAX-NORMING:

Even when looking at just the spikes, they assume a wide range of magnitudes. One potential method of dealing with this dynamic range is to scale tensors so that their max absolute element is 1 (similar to a per-tensor AdaMax (Kingma & Ba, 2014) or Range Batch-Norm (Banner et al., 2018) applied to gradients). Optimizers such as Adam, which normalize by gradient variance, provide a justification for why this sort of scaling might work well, although they work at a per-element rather than per-tensor level. We choose max-norming rather than variance-based norming because the former is easier computational and potentially more ammenable to quantization. However, a problem with the approach of normalizing tensors independently at each sample is that noise might be magnified during regions of quiet as seen in the Figure . What we therefore propose is normalization by the maximum of both the current max element and a moving average of the max element.

Explicitly, max-norm takes two parameters -a decay factor β = 0.999 and a gradient floor ε = 10 -4 and keeps two state variables -the number of evaluations k := 0 and the current maximum moving average x mv := ε. Then for a given input x, max-norm modifies its internal state and returns x norm : (Ioffe & Szegedy, 2015) normalizes a tensor X along some axes, then applies a trainable affine transformation. For each slice X of X that is normalized independently:

k := k + 1 x max := max(|x|) + ε x mv := β • x mv + (1 -β) • x max xmv := x mv 1 -β k x norm := x max(x max , xmv ) F STREAMING BATCH NORMALIZATION Standard batch normalization

Y = γ • X -µ b σ 2 b + ε + β

where µ b , σ b are mean and standard deviation statistics of a minibatch and γ, β are trainable affine transformation parameters.

In our case, we do not have the memory to hold a batch of samples at a time and must compute µ b , σ b in an online fashion. To see how this works, suppose we knew the statistics of each sample µ i , σ i for i = 1 . . . B in a batch of B samples. For simplicity, assume the i th sample is a vector X i,: ∈ R n containing elements X i,j . Then:

EQUATION

)

In other words, the batch variance is not equal to the average of the sample variances. However, if we keep track of the sum-of-square values of samples σ 2 i + µ 2 i , then we can compute σ 2 b as in (24). We keep track of two state variables: µ s , sq s which we update as µ s := µ s + µ i and sq s := sq s + σ 2 i + µ 2 i for each sample i. After B samples, we divide both state variables by B and apply (23, 24) to get the desired batch statistics. Unfortunately, in an online setting, all samples prior to the last one in a given batch will only see statistics generated from a portion of the batch, resulting in noisier estimates of µ b , σ b .

In streaming batch norm, we alter the above formula slightly. Notice that in online training, only the most recently viewed sample is used for training, so there is no reason to weight different samples of a given batch equally. Therefore we can use an exponential moving average instead of a true average to track µ s , sq s . Specifically, let:

µ s := η • µ s + (1 -η) • µ i sq s := η • sq s + (1 -η) • (σ 2 i + µ 2 i )

If we set η = 1 -1/B, a weighting of 1/B is seen on the current sample, just as in standard averages with a batch of size B, but now all samples receive similarly clean batch statistic estimates, not just the last few samples in a batch.

G ONLINE DATASET:

For our experiments, we construct a dataset comprising an offline training, validation, and test set, as well as an online training set. Specifically, we start with the standard MNIST dataset of LeCun et al. (1998) and split the 60k training images into partitions of size 9k, 1k, and 50k. Elastic transforms (Simard et al., 2003; Ernestus, 2016) are used to augment each of these partitions to 50k offline training samples, 10k offline validation samples, and 100k online training samples, respectively. Elastic transforms are also applied to the 10k MNIST test images to generate the offline test samples.

The source images for the 100k online training samples are randomly drawn with replacement, so there is a certain amount of data leakage in that an online algorithm may be graded on an image that has been generated from the same image a previous sample it has trained on has been generated from. This is intentional and is meant to mimic a real-life scenario where a deployed device is likely to see a restrictive and repetitive set of training samples. Our experiments include comparisons to standard SGD to show that SKS's improvement is not merely due to overfitting the source images.

From the online training set, we also generate a "distribution shift" dataset by applying unique additional augmentations to every contiguous 10k samples of the 100k online training samples. Four types of augmentations are explored. Class distribution clustering biases training samples belonging to similar classes to have similar indices. For example, the first thousand images may be primarily "0"s and "3"s, whereas the next thousand might have many "5"s. Spatial transforms rotate, scale, and shift images by random amounts. Background gradients both scale the contrast of the images and apply black-white gradients across the image. Finally, white noise is random Gaussian noise added to each pixel. In addition to distribution shift for testing adaptation, we also look at internal statistical shift of weights in two ways -analog and digital. For analog weight drift, we apply independent additive Gaussian noise to each weight every d = 10 steps with σ = σ 0 / 1M/d where σ 0 = 10 and re-clip the weights between -1 and 1. This can be interpreted as each cell having a Gaussian cumulative error with σ = σ 0 after 1M steps. For digital weight drift, we apply independent binary random flips to the weight matrix bits every d steps with probability p = p 0 /(1M/d) where p 0 = 10. This can be interpreted as each cell flipping an average of p 0 times over 1M steps. Note that in real life, σ 0 , p 0 depend on a host of issues such as the environmental conditions of the device (temperature, humidity, etc), as well as the rate of seeing training samples.

H HYPERPARAMETER SELECTION:

In order to compare standard SGD with the SKS approach, we sweep the learning rates of both to optimize accuracy. In Figure 6 , we compare accuracies across a range of learning rates for four different cases: SGD or SKS with or without max-norming gradients. Optimal accuracies are found when learning rate is around 0.01 for all cases. For most experiments, 8b weights, activations, and gradients, and 16b biases are used. Experiments similar to those in Section I are used to select some of the hyperparameters related to the SKS method in particular. In most experiments, rank-4 SKS with batch sizes of 10 (for convolution layers) or 100 (for fully-connected layers) are used. Additional details can be found in the supplemental code. 

I ADDITIONAL STUDIES:

In Figure 7 , rank and weight bitwidth is swept for SKS with gradient max-norming. As expected, training accuracy improves with both higher SKS rank and bitwidth. In dense NVM applications, higher bitwidths may be achievable, allowing for corresponding reductions in the SKS rank and therefore, reductions in the auxiliary memory requirements.

In Table 2 , biased (zero-variance) and unbiased (low-variance) versions of SKS are compared. Accuracy improvements are generally seen moving from biased to unbiased SKS although the pattern differs between the no-norm and max-norm cases. In the no-norm case, a significant improvement is seen favoring unbiased SKS for fully-connected layers. In the max-norm case, the choice of biased or unbiased SKS has only a minor impact on accuracy. It might be expected that as the number of accumulated samples for a given pseduobatch increases, lower variance would be increasingly important at the expense of bias. For our network, this implies convolutions, which receive updates at every pixel of an output feature map, would preferentially have biased SKS, while the fully-connected layer would preferentially be unbiased. This hypothesis is supported by the no-norm experiments, but not by the max-norm experiments.

In Table 3 , several ablations are performed on SKS with max-norm. Most notably, weight training is found to be extremely important for accuracy as bias-only training shows a ≈ 15 -30% accuracy hit depending on whether max-norming is used. Streaming batch norm is also found to be quite helpful, especially in the no-norm case. Now, we explain the κ th ablation. In Section 4.1.1, we found the SVD of a small matrix C and its singular values σ 1 , . . . , σ q . This allows us to easily find the condition number of C as κ(C) = σ 1 /σ q . We suspect high condition numbers provide relatively useless update information akin to noise, especially in the presence of L, R quantization. Therefore, we prefer not to update L, R on samples whose condition number exceeds threshold κ th . We can avoid performing an actual SVD (saving computation) by noting that C is often nearly diagonal, leading to the approximation κ(C) ≈ C 1,1 /C q,q Empirically, this rough heuristic works well to reduce computation load while having minor impact on accuracy. In Table 3 , κ th = 10 8 does not appear to ubiquitously improve on the default κ th = 100, despite being ≈ 2× slower to compute. 

:

This can alternately be achieved by sub-sampling the training data by r/B where r is the OK rank. The purpose of using a low-rank estimate is that for the same memory cost, it is significantly more informational than the sub-sampled data, allowing for faster training convergence.

Their target application differs slightly in that they handle matrix -vector Kronecker sums rather than vector -vector Kronecker sums.

The fact that it is unbiased: E[ Σ] = Σ can be easily verified.

As discussed in Appendix B.1, B < ni, so we substitute c in c/2||w tw * || with the minimum non-zero Eigenvalue of the Hessian c when plotting the RHS of (4).

Future work might look into how to change these clipping ranges, but this is beyond the scope of this paper.

SKS could potentially approximate Adam. SKS on a 2 , dz 2 , a, dz allows for a low-rank approximation of the variance of the gradients, however, this is unlikely to work well because of numerical stability (e.g., estimated variances might be negative).

The decision to use training data is deliberate, however experiments on out-of-sample images, such asRecht et al. (2019) show similar behavior.

Quantization ranges are chosen to optimize accuracy and are different from those in Section 7.1.

For example, one could compute just a sliding window of rows of every feature map, discarding earlier rows as later rows are computed, resulting in a square-root reduction of activation memory. To incorporate backpropagation, compute the forward pass once fully, then compute the forward pass again, as well as the backward pass using the sliding window approach in both directions.

Figures and Tables:

Figure1: In both plots, the solid line with markers plots the loss vs. gradient error variance (LHS of (4)) across 50 steps of SGD for several different setups. The left dashed line represents the RHS of (4) and the right dashed line is the RHS with C instead of c.

Figure 2: Adaptation of various training schemes over four different training environments (a) to (d). In each training environment, the top plot shows the exponential moving averages (0.999) of the per-sample online accuracy of the five training schemes, while the bottom plot shows the maximum number of updates applied to any given convolution or fully-connected kernel memory cell. For the distribution shifts in (b), the enabled augmentations at each contiguous 10k samples is shown (CD = class distribution, ST = spatial transforms, BG = background gradients, WN = white noise).

fully-connected layer transforms an input a ∈ R ni×1 to an intermediate z = W • a + b to an output y = σ(z) ∈ R no×1 where σ is a non-linear activation function. Gradients of the loss function with respect to the weight parameters can be found as:

Figure 4: Maximum magnitude of weight gradients versus training step for standard SGD on a CNN trained on MNIST.

Figure 5 shows some representative examples of what these augmentations look like. The augmentations are meant to mimic different external environments an edge devices might need to adapt to.

Figure 5: Samples of different types of distribution shift augmentations.

Figure6: The left two heat maps are used to select the base / standard SGD learning rate. The right two heat maps are used to select the SKS learning rate using the optimal SGD learning rate for bias training from the previous sweeps. For the SKS sweeps, the learning rate is scaled proportional to the square-root of the batch size B. This results in an approximately constant optimal learning rate across batch size, especially for the max-norm case. Accuracy is reported averaged over the last 500 samples from a 10k portion of the online training set, trained from scratch.

Figure 7: Accuracy across a variety of SKS ranks and weight bitwidths, showing the expected trends of increasing accuracy with rank and bitwidth. Accuracy is calculated by averaging the accuracy on the last 500 samples from a 2k portion of the training data. For bitwidths of 1 and 2, mid-rise quantization is used (e.g., 1 bit quantizes values to -0.5 and 0.5 instead of -1 and 0).

in the weight drift cases. Additionally, SKS shows about three orders of magnitude improvement compared to SGD in the worst case number of weight updates. Much of this reduction is due to the convolutions, where updates are applied at each pixel. However, reduction in fullyconnected writes is still important because of potential energy savings. SKS/max-norm performs best in terms of accuracy across all environments and has similar weight update cost to SKS/no-norm.7.2 TRANSFER LEARNING AND ALGORITHM COMPARISONSTo test the broader applicability of low rank training techniques, we run several experiments on ImageNet withResNet-34 (Deng et al., 2009;He et al., 2016), a potentially realistic target for dense NVM inference on-chip. For ImageNet-size images, updating the low-rank approximation at each pixel quickly becomes infeasible, both because of the single-threaded nature of the algorithm, and because of the increased variance of the estimate at larger batch sizes. Instead, we focus on training the final layer weights (1000 × 512). ResNet-34 weights are initialized to those fromPaszke et al. (2017) and the convolution layers are used to generate feature vectors for 10k ImageNet training images 7 , which are quantized and fed to a one-layer quantized 8 neural network. To speed up experiments, the layer weights are initialized to the pretrain weights, modulated by random noise that causes inference top-1 accuracy to fall to 52.7% ± 0.9%. In Table1, we see that the unbiased SKS has the strongest recovery accuracies, although biased SVD also does quite well. The high-variance UORO and true SGD have weak or non-existent recoveries. Accuracy recovery beyond inference (%, mean with standard deviation from 5 random seeds) between different algorithms (all with max-norm; effective batch size B = 100 if applicable), tested at different ranks (r), and learning rates (η). Optimal learning rates are bolded.

Table content: <table><tr><td>η 0.003</td><td>0.010</td><td>0.030</td><td>0.100</td><td>0.300</td></tr><tr><td>Algorithm r</td><td/><td/><td/><td/></tr></table>

Importance of unbiased SVD. Accuracy is calculated from the last 500 samples of 10k samples trained from scratch. Mean and unbiased standard deviation are calculated from five runs of different random seeds.

Table content: <table><tr><td colspan="2">Conv SKS FC SKS</td><td colspan="2">Accuracy (no-norm) Accuracy (max-norm)</td></tr><tr><td>Biased</td><td>Biased</td><td>79.7% ± 1.1%</td><td>82.7% ± 1.3%</td></tr><tr><td>Biased</td><td colspan="2">Unbiased 83.0% ± 0.9%</td><td>82.4% ± 1.2%</td></tr><tr><td>Unbiased</td><td>Biased</td><td>77.7% ± 1.5%</td><td>84.6% ± 2.0%</td></tr><tr><td>Unbiased</td><td colspan="2">Unbiased 81.0% ± 0.9%</td><td>83.6% ± 2.5%</td></tr></table>

Miscellaneous selected ablations. Accuracy is calculated from the last 500 samples of 10k samples trained from scratch. Mean and unbiased standard deviation are calculated from five runs of different random seeds.

Table content: <table><tr><td>Modified Condition</td><td colspan="2">Accuracy (no-norm) Accuracy (max-norm)</td></tr><tr><td colspan="2">baseline (no modifications) 80.2% ± 1.0%</td><td>83.0% ± 1.1%</td></tr><tr><td>bias-only training</td><td>51.8% ± 3.2%</td><td>68.6% ± 1.4%</td></tr><tr><td>no streaming batch norm</td><td>68.2% ± 1.9%</td><td>81.8% ± 1.3%</td></tr><tr><td>no bias training</td><td>81.3% ± 1.0%</td><td>83.0% ± 1.4%</td></tr><tr><td>κ th = 10 8 instead of 100</td><td>79.8% ± 1.4%</td><td>84.2% ± 1.4%</td></tr></table>

Peer review (iclr2020):
While inference on edge devices is a popular and well-studied problem in recent days, training on these devices comes with many challenges. This paper proposes a low-rank training schema that helps mitigate some of the critical challenges that occur during training models on NVM memory-based edge devices. Additionally, two techniques, namely streaming batch norm and gradient max norm, are proposed to help training in an online setting. The proposed method is mainly based on approximating the Kronecker sum and is largely inspired by (Benzing et al, ICML 2019, Optimal Kronecker-sum approximation of real time recurrent learning). The proposed approach provides a few optimizations that improves this performance further, and outperforms SGD in terms of accuracy and the number of weights updates in a limited experimental setting. +ves: + Training on an edge device is a relevant setting, where there is very little work so far, and this is a useful objective. + Focusing on the Kronecker sum and speeding it seems like an interesting solution to solving this problem. + The experimental settings considered (e.g. flipping bits and adding Gaussian noise to weights) is interesting, and perhaps of larger relevance to other work in the area. + The ablation studies provided in the appendix are useful and appreciable. Concerns: - The key concern is that most of the proposed method is built upon Benzing et al’s work (ICML 2019), and the original contribution seems limited. While the paper introduces a few optimizations further, this seems to be incremental than originally novel. - The problem is formulated for linear regression with least-squares loss, and the experimentation is carried out on the MNIST dataset (classification setting). How is the methodology relevant to the cross-entropy loss, typically used for classification? The paper does not talk about this. - The results are shown for a network with four 3 x 3 convolution layers and two fully connected layers on the MNIST dataset. Results with different architectures on a few other datasets (at least CIFAR-10) would be necessary to assess the usefulness of this method. Some discussion on what would be the maximum depth of the network that can be trained using this training schema and hardware would be very useful. - Do all deep neural network architectures (and loss functions) admit a Kronecker sum representation? What class of models can benefit from this method? Factorizations such as Cholesky allow to interpolate between computational complexity and decomposition rank by tuning a rank hyperparameter. Why would such factorizations not be better for memory-constrained settings? - It would be interesting to see the results using the proposed algorithm on standard hardware, will it provide the same performance as SGD when the scale of the problem increases in terms of layers and other regularization techniques, etc.? This would help understand the performance of the algorithm in a standard-setting. Minor issues: - The paper could have been organized better. Most of the main paper is used to describe Benzing et al’s method, and a lot of the details of the original contributions are in the Appendix. While the equations discussed in Section 4 hold for Linear Regression, the same cannot be directly extended for other networks like CNNs. Considering all the experiments are on CNNs, Sec B.2 in the Appendix should have been in the main paper to help a reader follow the experiments. =====POST-REBUTTAL COMMENTS======== I thank the authors for the rebuttal. The paper*s organization is much better now. Unfortunately, there are many concerns which are still not convincingly answered: (i) Considering the focus of this work (as admitted by the authors in the rebuttal) is empirical and with limited technical novelty, more comprehensive results on many datasets are required. To some extent, transfer to ImageNet addresses this partially, but not convincingly enough; (ii) It is not clear from the authors* rebuttal why testing this on standard hardware is not possible/appropriate. The rebuttal says *hard to quantify due to a variety of factors* but that is vague without stating the factors. I am willing to increase the rating from 3 to 4, but am unfortunately still towards rejecting the current shape of this work. I believe the work will benefit from considering all reviewers* comments more comprehensively to have a more impactful paper.