Towards Exact Gradient-based Training on Analog In-memory Computing

Given the high economic and environmental costs of using large vision or language models, analog in-memory accelerators present a promising solution for energy-efficient AI. While inference on analog accelerators has been studied recently, the training perspective is underexplored. Recent studies have shown that the"workhorse"of digital AI training - stochastic gradient descent (SGD) algorithm converges inexactly when applied to model training on non-ideal devices. This paper puts forth a theoretical foundation for gradient-based training on analog devices. We begin by characterizing the non-convergent issue of SGD, which is caused by the asymmetric updates on the analog devices. We then provide a lower bound of the asymptotic error to show that there is a fundamental performance limit of SGD-based analog training rather than an artifact of our analysis. To address this issue, we study a heuristic analog algorithm called Tiki-Taka that has recently exhibited superior empirical performance compared to SGD and rigorously show its ability to exactly converge to a critical point and hence eliminates the asymptotic error. The simulations verify the correctness of the analyses.

Paper

Similar papers

Peer review

Reviewer dVXd7/10 · confidence 3/52024-06-27

Summary

This is a theoretical paper about training analog systems based on resistive memories. The paper tackles specifically the problem of weight update asymmetry in such resistive devices. The paper develops a model for the weight dynamics under “Analog SGD”, taking into account weight update asymmetry, and it shows that Analog SGD does not converge, due to an asymptotic error. The paper then shows that the Tiki-Taka algorithm, developed in prior works for such analog devices, has better convergence properties. Tiki-Taka eliminates the asymptotic error and converges to a critical point. Simulations corroborate the theory.

Strengths

The topic tackled by the paper is very important. Analog computing could lead to energy efficiency gains of several orders of magnitude for both inference and training, compared to digital neural networks on GPUs. Developing appropriate theoretical models of analog computing to better understand the behavior of these analog systems is of utmost importance. To my knowledge, this paper is the first attempt to study the convergence properties of analog SGD. The paper is clear and the study is thorough.

Weaknesses

The theoretical analysis rests on a set of models and assumptions. Thus, the applicability of the conclusions rests on the validity of all these models/assumptions. In practice, I expect none of these assumptions to perfectly capture the real behavior of the system. For example, Figure 2 hints that the analog SGD dynamics (Eq3) better fits real behaviors than the digital SGD dynamics (Eq2), but the matching still seems to be far from perfect. This, of course, is not a problem, however, what is more problematic is that there does not seem to be any discussion in the paper about which assumptions/models are accurate and widely accepted, and which ones less accurately capture empirical data or are still debated by the community. The paper implicitly assumes that backpropagation (BP) is used for computing the weight gradients. However, a large (and growing) group of researchers in neuromorphic computing (or ”physical neural networks”) think that BP is not the best fit for analog devices, and they are exploring alternative training methods. See e.g. Ref [1] below for a very recent review of such algorithms. In particular, Ref [1] discusses several algorithms that extract or estimate the weight gradients in broad classes of analog systems. My understanding is that it wouldn’t take too much time and effort to extend the present study to include these other gradient-descent-based algorithms. In my view, including these other methods would add a lot of value to the paper, by broadening the potential impact of the work. Reference: [1] Momeni, Ali, et al. "Training of Physical Neural Networks." arXiv preprint arXiv:2406.03372 (2024). Minor remark. One claim in the introduction is not properly referenced. Specifically, where do the figures of $2.4 million and $4.6 million for training LLaMA and GPT-3 come from? The references provided are the LLaMA and GPT-3 papers, which do not provide these figures, if I am not mistaken.

Questions

Is the model of Eq6 for weight update asymmetry a widely accepted model? Does it match empirical data? Similarly, are the forms of q+ and q- for ALD, written at lines 174-175, widely accepted and/or corroborated by data? Does the AIHWKit very faithfully simulate real analog devices? What are the strengths and limitations of this tool? The paper is concerned with analog resistive memories where the weights are implemented as conductances (which are positive, I suppose), but on Figure 6 (Appendix L), I see that the weight values can be negative. Lemma 1 also indicates that w_k can be either positive or negative. Could you clarify this point? I understand that the RHS of Eq.5 can be arbitrarily small by choosing alpha = sqrt(1/K) and letting K -> infinity. But, unless I am mistaken, this does not imply that the sequence grad f(W_k) converges to zero. Can we conclude from Eq5 that W_k converges? If not, why is Eq.5 important for studying the convergence of (digital) SGD? I am not sure to understand why Theorem 2 is important. To my understanding, the important result is Theorem 3, which shows that analog SGD does not converge, contrary to digital SGD (Eq.5). Could you clarify this point? Is the theory of analog SGD and Tiki-Taka provided in this paper limited to backpropagation? Or could it be used with other methods that compute weight gradients, e.g. Equilibrium Propagation and/or other algorithms presented in Ref [1] above? Similarly, could the AIHWKit be employed with these other training algorithms?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The applicability of the theoretical results depends on the validity of the underlying assumptions – see my comment above. To be clear, I am supportive of the methodology followed by the authors. However, given that the overall model will be at most as accurate as the least accurate assumption, it would be extremely useful to know which assumptions/models we can trust 100% (i.e. are well accepted), and which ones are still debated / less supported by experiments. This would help future works identify what are the main “bottlenecks” of the study, to investigate those in depth and perhaps further improve the model. For instance, can we really trust Eq3 as “the right model” of weight dynamics? Or should one merely think of it as a better model than Eq. 2?

Reviewer gG9b6/10 · confidence 3/52024-07-11

Summary

The paper introduces the training on the analog in-memory accelerators with SGD. The traditional analog SGD algorithm suffers from inexact convergence due to asymmetric updates/gradient noise. The paper shows theoretical foundations for gradient-based training on analog devices. The authors propose Tiki-Taka, which reflects the device model to the Analog SGD. It shows better empirical performance, e.g., it can eliminate asymptotic errors and converge to critical points. They show the theoretical and empirical efficacy of the proposed approaches in overcoming training limitations on the analog devices.

Strengths

- The paper is well organized and easy to follow. - Concrete theoretical foundings on the proposed algorithm. - Benchmark sets and the target network look reasonable, and the result also sound and look promising. - Showing cases on various analog devices like ReRAM, etc.

Weaknesses

- It'd be great if authors present how computations (ops) can be realized on the analog devices. - Needs comparison to other analog SGD methods, like [16] or [22].

Questions

- Floating-point operations are not trivial on analog devices. How can they be implemented? Additionally, for SGD training, how are analog devices more efficient compared to digital devices? - How does it scale to large batch sizes?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Reviewer gG9b2024-08-09

Thanks for the detailed response. It has clarified some of my confusion and improved my understanding on the paper. For W2, to clarify my comment, it'd be great if you can elaborate differences to existing algorithms.

Reviewer G3ws5/10 · confidence 4/52024-07-16

Summary

This paper presents a theoretical framework for gradient-based training on analog devices. This work first identifies the non-convergence problem of Analog SGD, which stems from asymptotic errors due to asymmetric updates and gradient noise and then presents a convergence analysis of Tiki-Taka, demonstrating its ability to accurately converge to a critical point, thereby eliminating the asymptotic error.

Strengths

-- The proof the convergence of Analog SGD and showing that noise and asymmetric updates together lead to its asymptotic error. -- Demonstration of Tiki-Taka algorithm precisely converging to the critical point by mitigating the drift caused by asymmetric bias and noise. -- Empirical simulations using both synthetic and real datasets to confirm the presence of asymptotic error in Analog SGD and to show that Tiki-Taka outperforms Analog SGD.

Weaknesses

-- The experimental results were obtained on rather small datasets. It'd be great to include the results on ImageNet. -- The experiments were limited to one type of networks (i.e., resnet) on CIFAR dataset. How does it generalize on other networks such as Transformers, or mobile CNNs?

Questions

See weaknesses

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC