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

We really appreciate your recognizing the importance of our work and for helpful suggestions. Please find our point-by-point reply to your comments below.

W1. There is no discussion about which models are accurate and widely accepted. As the reviewer correctly pointed out, adding a discussion about the accuracy of the models are important. The applicability of our analysis relies on the following two analog hardware models:

• Asymmetric update model (6): The update (6) omits some non-ideality, like cycle-to-cycle/device-to-device variation and analog-digital conversion error [12][27]
• ALD response model (8): ALD models assume the response factor is a linear function, which is widely adopted in the literature [16,19,23].

It explains why the proposed dynamic does not match the real behaviors perfectly in Fig 2. We will clarify the potential error raised from the simplified model in the revision. We hope the discussion helps the readers better evaluate the scope of applicability of the theory.

W2. Including these other PNN training methods would add a lot of value, by broadening the potential impact of the work.

Thank you for pointing out this interesting direction and suggesting a recent work. Indeed, this is inspiring! To establish the connection between our work and the PNN training, we will add a paragraph to briefly review the progress as follows in literature review section.

We are especially interested in AIMC hardware with resistive crossbar arrays. It is a specific implementation of physical neural network (PNN) [R1,R2]. PNN is a model implemented by a physical system involving tunable analog quantity. The quantity is adjusted to implement the learning on some specific tasks. Various hardware is capable of supporting PNN, such as holographic grating [R3], wave-based systems [R4], and photonic networks [R5], to name a few.

[R1] Wright, et al. "Deep physical neural networks trained with backpropagation."

[R2] Momeni, et al. "Training of Physical Neural Networks."

[R3] Psaltis, et al. "Holography in artificial neural networks."

[R4] Hughes, et al. "Wave physics as an analog recurrent neural network."

[R5] Tait, et al. "Neuromorphic photonic networks using silicon photonic weight banks."

W3. Where do the figures of 4.6 million for training LLaMA and GPT-3 come from?

Thanks for the careful reading! The cost is estimated by multiplying the required GPU hours by the AWS price per hour, where the GPU hours are reported in [1]. To avoid confusion, we will rephrase this sentence with the following one.

For example, it requires 184 thousand GPU hours to train an LLAMA2 7 billion model [R6], and this time increases to 1.7 million GPU hours for its 70 billion version [1].

[R6] Touvron, et al. "Llama 2: Open foundation and fine-tuned chat models."

Q1. Is the model of Eq6 for weight update asymmetry a widely accepted model?

The answer is affirmative. For example, [16, 52] have already demonstrated different factors scale the update in different weight state.

Q2. Does the AIHWKit faithfully simulate real analog devices?

AIHWKit is capable of providing real simulations across different granular levels, including IO noise, pulse update, update variation, response factor, and A/D discretization, to name a few. One of the limitations is the significant overhand of detailed simulation and inadequate multi-GPU parallel support, which make it time-consuming to conduct large-scale simulations.

Q3. Why can the weights be negative? This is just for the mathematical convenience. To implement the negative weight on analog devices, two resistive crossbar arrays, a main array, and a reference array are set. The weight is represented by the difference of the two conductances multiplied by a scaling factor, which can be negative.

Thanks for pointing this out! We ignored this detail in our dynamic, but we will clarify it in the revision.

Q4. Can we conclude from Eq5 that $W_k$ converges? If not, why is Eq.5 important for studying the convergence of (digital) SGD?

In the convergence study of (digital) SGD, there are typically two types of convergence: the convergence to the stationary point (e.g., $W^*$ such that $\nabla f(W^*)=0$) or the convergence to the optimal solution (e.g., $W^*$ such that $W^*\in \arg\min_W f(W)$). Without additional assumptions like convexity, the convergence to the stationary point might be the best we can hope for in the worst case. Therefore, with a set of properly choosing decreasing stepsizes, Eq5 has been commonly used as a metric to assess the convergence to the stationary point of SGD; see [20]. As a result, we also use the metric of Eq5 to demonstrate the convergence of analog algorithms.

Q5. Why is Thm 2 important?

Compared to Thm 3, Thm 2 is important since it serves as another main component to reveal the performance limit of Analog SGD. Thm 3 claims that there exist a bad situation that Analog SGD has an asymptotic error but it does not ensure it is the worst case. Combining Thm 2 and 3, we claim that $4\sigma^2 S_K$ is the worst-case possible asympotic error, which provide a better insight.

Q6. Is the theory provided in this paper limited to backpropagation?

Thanks for the inspiring question! We believe our theory can be adapted to other analog training algorithms since algorithms like equilibrium propagation (EP) adopt other methods to determine the update directions. To replace the gradient in our analysis with a given update direction, we can still study the convergence using similar techniques. We will seriously consider this as a future direction.

Q7. Could the AIHWKit be employed with these other training algorithms?

The answer is affirmative. AIHWKit enables multiple levels of simulation granularity. To implement other algorithms like EP, one needs to modify the algorithm-level code.

With the above clarifications, we hope that our responses addressed all your questions.

We thank the reviewer for acknowledging the merits of our work. Our point-to-point response to your comments and suggestions follow next.

W1. It'd be great if authors present how computations (ops) can be realized on the analog devices.

The workload of forward and backward computation can be separated into two categories: matrix-vector multiplication (MVM) operations and non-MVM operations. In AIM computation hardware, the MVM operations are implemented in analog in-memory arrays, while the other non-MVM operations are conducted in digital circuits.

We illustrate in the attached PDF in Author Rebuttal on how the MVM is implemented by analog devices. In a $3\times 3$ resistive crossbar array, a matrix $W$ is represented by the conductance of resistor at the crosspoint, where the $(i,j)$-th element of $W$ is represented by the conductance of the $(i,j)$-th resistor. To conduct an MVM operation $z = Wx$, voltage $x_{j}$ is applied between $j$-th and $(j+1)$-th row, where the subscript $j$ is used to imply the $j$-th coordinate. By Ohm's law, the current is $I_{ij}=W_{ij}x_{j}$; and by Kirchhoff's law, the total current on the $i$-th column is $\sum_{j}I_{ij}=\sum_{j}W_{ij}x_{j}$.

W2. Needs comparison to other analog SGD methods, like [16] or [22].

Note that the goal of this paper is not to develop the state-of-the-art analog training algorithm, but to understand why the vanilla SGD training on analog device does not work as expected, and why some correction operations used in the heuristic algorithms such as Tiki-Take work, through the unified lens of asymmetric updates on analog devices. We hope this theoretical understanding can contribute to future developments of analog training algorithms, hardware and material.

In this context, in the introduction, we have compared the studied algorithm Tiki-Taka (TT-v1) with TT-v2 [22] and TT-v3/v4 [16] (c.f. Sec 1.2). Both [16] and [22] are variants of Tiki-Taka which were proposed to deal with some practical issues in analog training such as the reading noise and non-perfect zero-shifting issues. Since our focus in this paper is on understanding the impact of asymmetric updates on analog devices, in the simulations parts, we omitted the comparison between our paper and [16][22].

Nevertheless, as requested by the reviewer, we have now compared these four methods in the same setting as that in Section 5.3. We train a CNN model under $\tau=0.7$ on MNIST dataset, whose results are listed below. TT-v2--v4 are always better than TT-v1.

Q1. How can floating-point (FP) be implemented by analog devices?

This is a great question! The FP is implemented in the digital domain [12, 31]. In the MVM computations, the weights are represented by conductance, while the input and output are represented by current signals. Therefore, the FP is unnecessary for MVM operations. The other operations involving FP computation are non-MVM ones, which are conducted in digital domains. We will add a remark in the revision to clarify it.

Q2. How does it scale to large batch sizes?

This is an interesting point! Implementing mini-batch gradient computation on the analog devices are significantly different from that on the digital devices. To implement the large batch in AIM computation hardware, each gradient in a batch is computed sequentially and accumulated to $W_k$ (for Analog SGD) or $P_k$ (for Tiki-Taka). For example, we use batch size 8 in FCN/CNN training and 128 for Resnet training. Studying the impact of batch size on the convergence of analog training will be an interesting future direction!

We hope the above detailed clarifications can fully resolve your concerns.

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.

We thank the reviewer for the time spent in reviewing our paper and for the valuable comments.

The weaknesses identified by the reviewer mainly focus on the datasets and the model archtecture used in the experimental results as follows.

• 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?

Albeit the main focus of this paper is on the theoretical understanding of analog training, we have still conducted additional experiments during the limited rebuttal period.

To demonstrate the efficiency of Tiki-Taka under various situations, we conducted more simulations on different datasets and model architectures, including MobileNetV2 and large/small MobileNetV3 on CIFAR10/CIRAR100 datasets. The results are listed as follows.

Table 1: Training on CIFAR10 dataset

Table 2: Training on CIFAR100 dataset

The results show that Tiki-Taka always achieves better test accuracy than Analog SGD by about 0.5% in almost all cases, which is consistent with our conclusion. Due to the limited time and inadequate GPU resources, we can not perform simulations on a larger scale like training on ImageNet dataset during the rebuttal period.