Year: 2019

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.

text: 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.

heading: 1 INTRODUCTION
text: 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 (Konecný 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 106 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 TRAINING NON-VOLATILE MEMORY
text: The meat of most deep learning systems are many weight matrix - activation vector productsW · 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 um2 (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 no × ni, we should use at most r(ni + no)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.

heading: 4 LOW-RANK TRAINING METHOD
text: Let z(i) = Wa(i) + b be the standard affine transformation building block of some larger network, e.g., y(i)p = fpost(z(i)) and a(i) = fpre(x(i)) with prediction loss L(y(i)p ,y(i)t ), where (x(i),y (i) t )
is the ith training sample pair. Then weight gradient∇WL(i) = dz(i) ( a(i) )> = dz(i) ⊗ a(i) where dz(i) = ∇z(i)L(i). A minibatch SGD weight update accumulates this gradient over B samples: ∆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̃ ∈ Rno×r, R̃ ∈ Rni×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.

heading: 4.1 OPTIMAL KRONECKER SUM APPROXIMATION (OK)
text: 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)] 6= X . Benzing et al. (2019) solves these problems by proposing a minimum variance unbiased estimator for rankReduce, which they call the OK algorithm2.
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).

heading: 4.1.1 EFFICIENT SVD OF KRONECKER SUMS
text: LetL = [L̃,dz(i)] andR = [R̃,a(i)] so thatLR> = L̃R̃>+dz(i)⊗a(i). Recall that rankReduce should turn rank-q LR> into an updated rank-r L̃R̃>.
QR-factorize L = QLRL and R = QRRR where QL ∈ Rno×q,QR ∈ Rni×q are orthogonal so that LR> = QL(RLR>R)Q > R. Let C = RLR > R ∈ Rq×q. Then we can find the SVD of C = UCΣV > C in O(q3) time (Cline & Dhillon, 2006), making it computationally feasible on small devices. Now we have:
1This 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.
2Their target application differs slightly in that they handle matrix - vector Kronecker sums rather than vector - vector Kronecker sums.
LR> = QL(UCΣV > C )Q > R = (QLUC)Σ(QRVC) > (1)
which gives the SVD of LR> sinceQLUC andQRVC are orthogonal and Σ is diagonal. This SVD computation has a time complexity ofO((ni+no+q)q2) and a space complexity ofO((ni+no+q)q).

heading: 4.1.2 MINIMUM VARIANCE, UNBIASED ESTIMATE OF Σ
text: 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:
m = min i s.t. (q − i)σi ≤ q∑
j=i
σj k = q −m
x0 =
(√ 1− σmk
s1 , . . . ,
√ 1− σqk
s1
)> s1 =
q∑ i=m σi
Note that ||x0||2 = 1. LetX ∈ R(k+1)×(k) be orthogonal such that its left nullspace is the span of x0. ThenXX> = I − x0x>0 . Now, let s ∈ {−1, 1}(k+1)×1 be uniform random signs and define:
Xs = (s X:,1, . . . , s X:,k) Z = √ s1 k ·Xs
Σ̃L = Σ̃R = diag (√ σ1, . . . , √ σm−1,Z ) (2)
where is an element-wise product. Then Σ̃LΣ̃>R = Σ̃ is a minimum variance, unbiased3 rank-r approximation of Σ. Plugging Σ̃ into (1),
LR> = (QLUC)Σ(QRVC) > ≈ (QLUC)Σ̃(QRVC)> = (QLUCΣ̃L)(QRVCΣ̃R)> (3)
Thus, L̃ = QLUCΣ̃L ∈ Rno×r and R̃ = QRVCΣ̃R ∈ Rni×r gives us a minimum variance, unbiased, rank-r approximation L̃R̃>.

heading: 4.2 STREAMING KRONECKER SUM APPROXIMATION (SKS)
text: 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 QL,QR
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 QL,QR, and weightings cx such that L̃ = QL · diag( √ cx)[:r] and R̃ = QR · diag( √ cx)[:r]. Upon receiving a new sample, a single inner loop of the numerically-stable modified Gram-Schmidt (MGS) algorithm (Björck, 1967) can be used to update QL and QR. The orthogonal basis coefficients cL = Q>Ldz (i) and cR = Q>Ra (i) computed during MGS can be used to find the new value of C = cLc>R + diag(cx).
3The fact that it is unbiased: E[Σ̃] = Σ can be easily verified.
After computing Σ̃L = Σ̃R in (2), we can orthogonalize these matrices into Σ̃L = Σ̃R = QxRx. Then from (3), we have L̃R̃> = (QLUCQx)(RxR>x )(QRVCQx)
>. With this formulation, we can maintain orthogonality inQL,QR by setting:
QL ← QLUCQx QR ← QRVCQx cx ← diag(RxR>x )
These matrix multiplies requireO((ni+no)q2) 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 (RxR > x ) is diagonal.

heading: 4.2.2 ORTHOGONALIZATION OF Σ̃L
text: 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 matrixRx with diagonal elements √ cx:
Qx =
[ Im−1 0
0 Xs
] √ cx = ( √ σ1, . . . , √ σm−1, √ s1/k︸ ︷︷ ︸
q−m+1 times
)

heading: 4.2.3 FINDING ORTHONORMAL BASIS X
text: We generated X by finding an orthonormal basis that was orthogonal to a vector x0 so that we could haveXX> = I − x0x>0 . An efficient method of producing this basis is through Householder matrices (x0,X) = I − 2 vv>/||v||2 where v = x0 − e(1) and (x0,X) is a k + 1× k + 1 matrix with first column x0 and remaining columnsX (Householder, 1958; user1551, 2013).

heading: 4.2.4 EFFICIENCY COMPARISONS TO STANDARD APPROACH
text: The OK/SKS methods require O((ni + no + q)q2) operations per sample and O(ninoq) operations after collecting B samples, giving an amortized cost of O((ni + no + q)q2 + ninoq/B) operations per sample. Meanwhile, a standard approach expands the Kronecker sum at each sample, costing O(nino) operations per sample. If q B,ni, no then the low rank method is superior to minibatch SGD in both memory and computational cost.

heading: 5 CONVEX CONVERGENCE
text: 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(wt) for flattened weight vectorwt and online sample t. In Appendix B, we show that with inverse squareroot learning rate, when the loss landscape Hessians satisfy 0 ≺ cI ∇2f t(wt) 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 ||wt −w∗|| (4) R(T ) = T∑ t=1 f t(wt)− T∑ t=1 f t(w∗) (5)
B∑ i=1 ( σ(t,i)q )2 ≤ c 2 4 ||wt −w∗||2 (6)
B∑ i=1 σ(t,i)r σ (t,i) q ≤ c2 8 ||wt −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.

heading: 5.1 CONVERGENCE EXPERIMENTS
text: We validate (4) with several linear regression experiments on a static input batch X ∈ R1024×100 and target Yt ∈ R256×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.

heading: 6 IMPLEMENTATION DETAILS
text: 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 sample6. 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.
4As discussed in Appendix B.1, B < ni, so we substitute c in c/2||wt −w∗|| with the minimum non-zero Eigenvalue of the Hessian c̃ when plotting the RHS of (4).
5Future work might look into how to change these clipping ranges, but this is beyond the scope of this paper. 6SKS could potentially approximate Adam. SKS on a2,dz2,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).

heading: 7 EXPERIMENTS

heading: 7.1 ADAPTATION EXPERIMENTS
text: 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. First, a control case where both external/environment and internal/NVM drift statistics are exactly the same as during offline training. Second, a case where the input image statistical distribution shifts every 10k samples, selecting from augmentations such as spatial transforms and background gradients (see Section G). Third and fourth are cases where the NVM drifts from the programmed values, roughly modeling NVM memory degradation. In the third case, Gaussian noise is applied to the weights as if each weight was a single multi-level memory cell whose analog value drifted in a Brownian way. In the fourth case, random bit flips are applied as if each weight was represented by b memory cells (see Appendix G for details). For each hypothetical scenario, we plot five different training schemes: pure quantized inference (no training), bias-only training, standard SGD training, SKS training, and SKS training with max-normed gradients. In SGD training and for training biases, parameters are updated at every step in an online fashion. These are seen as different colored curves in Figure 2.
Inference does best in the control case, but does poorly in adaptation experiments. SGD doesn’t improve significantly on bias-only training, likely because SGD cannot accumulate gradients less than a weight LSB. SKS, on the other hand, shows significant improvement, especially after several
thousand samples 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.

heading: 7.2 TRANSFER LEARNING AND ALGORITHM COMPARISONS
text: To test the broader applicability of low rank training techniques, we run several experiments on ImageNet with ResNet-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 from Paszke et al. (2017) and the convolution layers are used to generate feature vectors for 10k ImageNet training images7, which are quantized and fed to a one-layer quantized8 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 Table 1, 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.

heading: 8 CONCLUSION
text: 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.
7The decision to use training data is deliberate, however experiments on out-of-sample images, such as Recht et al. (2019) show similar behavior.
8Quantization ranges are chosen to optimize accuracy and are different from those in Section 7.1.

heading: A SKS ALGORITHM
text: Algorithm 1 Streaming Kronecker Sum Approximation
State: QL ∈ Rno×q;QR ∈ Rni×q; cx ∈ Rq×1 Input: dz(i) ∈ Rno×1; a(i) ∈ Rni×1 for i ∈ [1, B] for i = 1 . . . B do {Modified Gram-Schmidt.} cL, cR ← 0q×1 for j = 1 . . . r do cL,j ← QL,j · dz(i); dz(i) ← dz(i) − cL,j ·QL,j cR,j ← QR,j · a(i); a(i) ← a(i) − cL,j ·QL,j
end for cL,q ← ||dz(i)||; QL,q ← dz(i)/cL,q cR,q ← ||a(i)||; QR,q ← a(i)/cR,q {Generate C and find its SVD.} C ← cLc>R + diag(cx) UC · diag(σ) · V >C ← SVD(C) {Minimum-variance unbiased estimator for Σ.} m← min j s.t. (q − j)σj ≤ ∑q `=j σ`
s1 ← q∑
i=m
σi, k ← q −m v ← √
1− k/s1 · σ[m:] − e(1)
s← {−1, 1}(k+1)×1 {Ind. uniform random signs.} Xs ← ( I + (s v)(v/v1)> ) [2:] {Householder.} {QR-factorization of Σ̃L.}
Qx ← [ I 0 0 Xs ] ∈ Rq×r cx ← (σ1, . . . , σm−1, s1/k, . . . , s1/k︸ ︷︷ ︸ q−m+1 times ) {Update the first r columns ofQL,QR.} QL[:r] ← QL ·UC ·Qx QR[:r] ← QR · VC ·Qx
end for {Compute final L̃, R̃ where∇WL ≈ L̃R̃>.} L̃← ( QL · diag( √ cx) ) [:r]
R̃← ( QR · diag( √ cx) ) [:r]

heading: B CONVEX CONVERGENCE
text: 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 = ni · no). 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 = maxw,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(wt) that is strongly convex in weight parameters wt, so there are positive constants C ≥ c > 0 such that cI ∇2f t(wt) CI for all t (Boyd & Vandenberghe, 2004, Section 9.3). We define regret as R(T ) = ∑T t=1 f t(wt) − ∑T t=1 f t(w∗)
where w∗ = argminw ∑T t=1 f t(w) (i.e., it is an optimal offline minimizer of f1, · · · , fT ).
The gradients seen during SGD are gt = ∇f t(wt) and we assume they are bounded by G = maxw∈F,t∈[1,T ]||∇f t(w)||. We also assume errors are bounded by E = maxt∈[1,T ]||εt||. Therefore, maxt∈[1,T ]||g̃t|| ≤ maxt∈[1,T ]||gt||+ ||ε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 gt · (wt −w∗)− c2 ||w t −w∗||22 ≤ g̃t · (wt −w∗), we have sublinear regret:
R(T ) ≤ D 2
2
√ T + (G+ E)2 (√ T − 1/2 ) Proof. From strong convexity cI ∇2f t(wt) CI for all t,
f t(w) + gt · (v −w) + c 2 ||v −w||22 ≤ f t(v) for all v (8)
In particular, if we consider v = w∗ and rearrange,
f t(w)− f t(w∗) ≤ gt · (w −w∗)− c 2 ||w −w∗||22 f t(w)− f t(w∗) ≤ g̃t · (wt −w∗) (9)
Consider a gradient update wt+1 = PF(wt − ηtg̃t), where PF projects the update back to F. Then,
||wt+1 −w∗||22 = ||P (wt − ηtg̃t)−w∗||22 ≤ ||wt − ηtg̃t −w∗||22 = ||wt −w∗||22 − 2ηt(wt −w∗) · g̃t + η2t ||g̃t||22 ≤ ||wt −w∗||22 − 2ηt(wt −w∗) · g̃t + η2t (G+ E)2
g̃t · (wt −w∗) ≤ 1 2ηt
( ||wt −w∗||22 − ||wt+1 −w∗||22 ) + ηt 2 (G+ E)2 (10)
From (9, 10),
f t(wt)− f t(w∗) ≤ 1 2ηt
( ||wt −w∗||22 − ||wt+1 −w∗||22 ) + ηt 2 (G+ E)2 (11)
We now bound the regret:
R(T ) = T∑ t=1 [ f t(wt)− f t(w∗) ] ≤
T∑ t=1 [ 1 2ηt ( ||wt −w∗||22 − ||wt+1 −w∗||22 ) + ηt 2 (G+ E)2 ]
= ||w1 −w∗||22 2η1 − ||w T+1 −w∗||22 2ηT + 1 2 T∑ t=2 ( 1 ηt − 1 ηt−1 ) ||wt −w∗||22 + (G+ E)2 2 T∑ t=1 ηt
≤ ||w 1 −w∗||22
2η1 +
1
2 T∑ t=2 ( 1 ηt − 1 ηt−1 ) ||wt −w∗||22 + (G+ E)2 2 T∑ t=1 ηt
≤ D 2
2η1 +
1
2 T∑ t=2 ( 1 ηt − 1 ηt−1 ) D2 + (G+ E)2 2 T∑ t=1 ηt
= D2
2ηT +
(G+ E)2
2
T∑ t=1 ηt (12)
If ηt = 1/ √ t, then ∑T t=1 ηt ≤ 2 √ T − 1 (Zinkevich, 2003), so from (12),
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 gt · (wt −w∗)− c2 ||w
t −w∗||22 ≤ g̃t · (wt −w∗). We now examine sufficient conditions for this inequality to be satisfied.
gt · (wt −w∗)− c 2 ||wt −w∗||22 ≤ g̃t · (wt −w∗)
εt · (w∗ −wt) ≤ c 2 ||wt −w∗||22 (14)
Since εt · (w∗ −wt) ≤ ||εt|| · ||w∗ −wt|| by Cauchy-Schwarz, it is sufficient for:
||εt|| · ||w∗ −wt|| ≤ c 2 ||wt −w∗||22
||εt|| ≤ c 2 ||wt −w∗|| (15)
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 < ni then c = 0. We can modify (15) to handle this case. Let c̃ be the minimum non-zero Eigenvalue of XX> and let w̃ ∈ RB represent w ∈ RN in the Eigenbasis of XX>. Then (15) becomes:
||εt|| ≤ c̃ 2 ||w̃t − w̃∗|| (16)
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 s21/k + s2 where s1 = ∑q i=m σi, s2 = ∑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, s1 = σr + σq, and s2 = σ2r + σ 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 ||wt −w∗||2 (19)
B.3 DISCUSSION ON CONVERGENCE
Equation (19) suggests that as wt → 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:
B∑ i=1 ( σ(t,i)q )2 ≈ B∑ i=q ( σi(G t) )2 ≤ B∑ i=1 ( σi(G t) )2 = ||Gt||2F
whereGt = ∇W tf t(W t) ∈ R(no×ni) are the matrix weightW t gradients at batch t and || · ||F is a Frobenius norm. We therefore expect both the left (proportional to ||Gt||2F ) and the right (proportional to ||wt −w∗||2) of (19) to decrease during training as wt → 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.
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:
N ∆2
12 + B∑ i=1 ( σ(t,i)q )2 ≤ c 2 4 ||wt −w∗||2 (20)

heading: C KRONECKER SUMS IN NEURAL NETWORK LAYERS
text: C.1 DENSE LAYER
A dense or fully-connected layer transforms an input a ∈ Rni×1 to an intermediate z = W · a+ b to an output y = σ(z) ∈ Rno×1 where σ is a non-linear activation function. Gradients of the loss function with respect to the weight parameters can be found as:
∇WL = (∇zL)︸ ︷︷ ︸ 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 ∈ Rhin×win×cin to an intermediate feature map Z = Wkern ∗ A + b ∈ Rhout×wout×cout through a 2D convolution ∗ with weight kernel Wkern ∈ Rcout×kh×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 matrixAcol ∈ R(houtwout)×(khkwcin) where the ith row is a flattened version of the sub-tensor of a which is dotted with Wkern to produce the ith pixel of the output feature map (Ren & Xu, 2015). We can multiply Acol by a flattened version of the kernel, W ∈ Rcout×(khhwcin) to perform the Wkern ∗ A convolution operation with a matrix multiplication. Under the matrix multiplication interpretation, weight gradients can be represented as:
∇WL = (∇ZcolL)︸ ︷︷ ︸ dZ>col (∇WZ)︸ ︷︷ ︸ Acol =
houtwout∑ i=1 dZ>col,i ⊗A>col,i (22)
which is the same as houtwout 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 Bhoutwout/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.

heading: D HARDWARE QUANTIZATION MODEL
text: 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 thatW `, 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.
9For 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.
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).

heading: E GRADIENT MAX-NORMING
text: 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. 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 xmv := ε. Then for a given input x, max-norm modifies its internal state and returns xnorm:
k := k + 1
xmax := max(|x|) + ε xmv := β · xmv + (1− β) · xmax x̃mv :=
xmv 1− βk
xnorm := x
max(xmax, x̃mv)

heading: F STREAMING BATCH NORMALIZATION
text: Standard batch normalization (Ioffe & Szegedy, 2015) normalizes a tensor X along some axes, then applies a trainable affine transformation. For each sliceX of X that is normalized independently:
Y = γ · X − µb√ σ2b + ε + β
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 ith sample is a vectorXi,: ∈ Rn containing elements Xi,j . Then:
µb = 1
B B∑ i=1 µi (23)
σ2b = 1
B B∑ i=1 1 n n∑ j=1 X2i,j − µ2b = 1 B B∑ i=1 ( σ2i + µ 2 i ) − µ2b 6= 1 B B∑ i=1 σ2i (24)
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 σ2i + µ 2 i , then we can compute σ 2 b as in (24). We keep track of two state variables: µs, sqs which we update as µs := µs+µi and sqs := sqs+σ2i +µ 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, sqs. Specifically, let:
µs := η · µs + (1− η) · µi sqs := η · sqs + (1− η) · (σ2i + µ2i )
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.

heading: G ONLINE DATASET
text: 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. 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.
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 = p0/(1M/d) where p0 = 10. This can be interpreted as each cell flipping an average of p0 times over 1M steps. Note that in real life, σ0, p0 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.

heading: H HYPERPARAMETER SELECTION
text: 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.

heading: I ADDITIONAL STUDIES
text: 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) ≈ C1,1/Cq,q. Empirically, this rough heuristic works well to reduce computation load while having minor impact on accuracy. In Table 3, κth = 108 does not appear to ubiquitously improve on the default κth = 100, despite being ≈ 2× slower to compute.

Review summaries:
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.

This paper proposes a low rank training method called the Streaming Kronecker Sum approximation (SKS algorithm) for training low precision models on edge devices. The authors compare their method to SGD for convolutional networks on MNIST and demonstrate improvements in terms of accuracy. The authors make use of the Optimal Kronecker-sum algorithm of Benzing et al and propose further improvements to it in the form of the SKS algorithm.