Analog In-memory Training on General Non-ideal Resistive Elements: The Impact of Response Functions

Thanks for acknowledging the theoretical contribution of our work and the support. We address other comments as follows.

W1: The experimental aspects of the work can be improved to make the proposed work more convincing. Specifically, the application of the proposed residual learning algorithm and its variants have been only demonstrated on fairly small datasets for which, some modest performance gains have been observed compared to other recent methods.

Thanks for your valuable suggestion! We incorporate more simulations to make the paper more convincing, including training language models from scratch and their finetuning. Compared with the existing simulations, the additional ones involve deeper models and have more weights. The simulation results, presented in the response to Q2 of reviewer GpyD, show that the proposed approaches generalize well on more challenging tasks.

---

W2: More details of physical modeling of AMIC systems and weight updating process can be provided. For example, it is not clear what kind of I/O reading errors have been considered in the experimental studies.

Thanks for your valuable suggestion. In all simulations, we adopt the following I/O settings. We will attach the details in the revision.

Table T4.1: I/O setting in the paper

---

Q1: While the proposed methods demonstrate performance improvements over the vanilla analog SGD, no discussion has been provided regarding its computational overhead. Is there any significant increase of computational complexity? How does it scale with the model size?

Thanks for your construction question! We compare the computational overhead in the response to Q4 in reviewer GpyD. Although an extra residual array $P_k$ is introduced, the latency does not significantly increase. Furthermore, as existing work presents, the resistive crossbar array implements the MVM operator with $O(1)$ time complexity as the matrix size increases, ensuring its strong scalability [R1.5].

---

We hope our response resolves your concerns. We are keen to engage in further discussion to ensure the clarity and rigor of our work. Thanks!

Thanks for acknowledging the theoretical contribution of our work and the support. We address some other comments as follows.

Q1: I wonder if it would be possible to incorporate more simulations with more practical scenarios and other realistic imperfections in the rebuttal period.

Thanks for your constructive suggestion! Following your suggestion, we incorporated more simulations to reflect more practical setting challenges in analog training:

1) More challenging machine learning tasks. We conducted a series of additional simulations on various tasks, including training language models from scratch and their finetuning. Compared with the existing simulation results, the additional ones involve deeper models and more parameters in the analog domain. The simulation results, presented in the response to Q2 in reviewer GpyD, show that the proposed approaches generalize well on more challenging tasks.

2) Validation on real resistive element configuration. To validate the approaches in real resistive elements, we performed simulations on various resistive elements, including ReRAM, EcRAM, and EcRAM-MO. The training on them is more challenging since they suffer from the asymmetric update issue, and they have limited update granularity and dynamic ranges. The simulation results, presented in the response Q1 in reviewer Kwur, show that it is promising to apply the proposed methods on real resistive elements.

We hope our response resolves your concerns. We are keen to engage in further discussion to ensure the clarity and rigor of our work. Thanks!

Thanks for acknowledging that the identified problem in this paper is crucial to overcome for successful training on AIMC hardware. We address other comments as follows.

---

W1 & W3 & Q1. No evaluation on any actual analog hardware. No evaluation on a specific simulated analog hardware. How would you validate your approach on either a hardware model or actual hardware?

Thanks for your constructive suggestion! As mentioned in the response to Q5 of Reviewer GpyD, our work builds upon a wide range of empirical study [45-59] and all our simulations are done on AIHWKIT simulator [44]. Since the AIHWKIT already captures most of the known hardware imperfections, the conclusion from our paper should be able to generalize to actual hardware.

Although we do not perform it in this work, we will discuss how to validate the proposed approach on actual hardware. We will follow the methods presented by [R2.1], which proposes a device-algorithm co-optimization framework that validates the Tiki-Taka [22-23] on HfO$_2$-based ReRAM with 14nm CMOS technology. Similar to them, we will first need to measure the response curve of a given resistive element (e.g., ReRAM) and simulate its behaviour in the analog hardware simulator. We then validate that the proposed approach can achieve comparable accuracy with digital training benchmarks on different machine learning tasks. Based on the simulation results, we can optimize the hardware to reproduce the simulation results on actual hardware.

As the first step, we validate the proposed approaches on two real resistive element configurations. The results presented in Table T2.1 suggest that the proposed algorithm works well on various real resistive elements. The rest of the hardware validation involves chip design, fabrication, and extensive testing. We will do so in future work.

Table T2.1: Test accuracy of training FCN models on MNIST dataset

---

W2. Limited evaluation on more larger machine learning datasets like Imagenet.

Thanks for your valuable suggestion! We include additional simulations in the revision to make our paper more convincing. In the response to Q2 of the reviewer GpyD, we present additional simulation results in the language model training during the limited rebuttal period, including training from scratch and finetuning tasks. Since the simulations on ImageNet training take a longer time, we will complete them after the rebuttal and add them in the revision.

---

W4 & Q2. The paper would benefit from an illustration of the proposed alternative algorithm on toy datasets that meet the assumptions of the theoretical results.

Thanks for your constructive suggestions! We will add a new subsection to discuss and illustrate the algorithms at the end of Section 4. We consider a toy regression problem $\min_{W\in{R}^2}||W-W^*||^2$, whose objective is smooth and strongly convex.

We choose a symmetric point $W^\diamond\ne W^*$. We add a Gaussian noise with mean 0 and variance $\sigma^2\ne 0$ to the gradients. We plot the convergence trajectories of Analog SGD (equation (2)) and Residual Learning (equation (10)-(11)).

We show that the weights generated by Analog SGD are attracted toward the symmetric point $W^\diamond$, as suggested by Theorem 1. Instead, Residual Learning converges to the optimal point and the auxiliary array $P_k$ is a good approximation of the residual $\frac{W^*-W_k}{\gamma}$, as suggested by the theory. We will discuss this example in the revision to provide more insight.

---

We hope our response resolves the concerns. Thanks!

---

[R2.1] Gong, et al. Deep learning acceleration in 14nm CMOS compatible ReRAM array: device, material and algorithm co-optimization. IEDM 2022.

[R2.2] Kim, et al. Metal-oxide based, CMOS-compatible ECRAM for deep learning accelerator. IEDM 2019.

Thanks for acknowledging that this paper provides a strong theoretical foundation with rigorous mathematical analysis, and the support. We address your comments as follows.

---

Q1: The experimental scope is somewhat limited, as simulations are conducted using AIHWKIT rather than actual hardware, which constrains real-world validation.

Thanks for your valuable suggestion! Analog in-memory computing is an emerging technology that attracts researchers from material science, device design, chip design, computer systems, and AI algorithm design. Our focus was mainly on the algorithm design for analog in-memory computing, where evaluation using off-the-shelf hardware simulators is very common. Since the reliable devices for analog in-memory training are not commercially available, we plan to leave the hardware evaluation for future work. Nonetheless, we still perform additional simulations to validate the approaches in real resistive element configurations, including ReRAM, EcRAM, and EcRAM-MO. The simulation results, presented in the response to Q1 of reviewer Kwur, demonstrate the potential in real-world applications.

---

Q2. The dataset scale remains relatively modest with MNIST and CIFAR experiments, raising questions about scalability to larger, modern datasets (however, most papers on analog training adopt these small-scale datasets as benchmarks)

Thanks for your valuable suggestion. We evaluate the scalability by investigating their performance on more training tasks, where the models are deeper or have more weights. The preliminary results below indicate that analog training is promising for scaling to larger and more modern datasets and tasks.

Task 1: Training language models from scratch. We train a Transformer model from scratch on the Coriolanus dataset, a play by Shakespeare. The model consists of 6 layers with ~10 million parameters. The weights for all query, key, value, and output projections in attention layers, as well as linear layers in feed-forward networks, are updated in the analog domain. The other parameters, such as token/position embeddings, scale/bias in layer normalization layers, are put and updated in the digital domain. For a fair comparison, we do not use an adaptive learning rate for either digital or analog training. In Table T1.1, we reach the same conclusion that Residual Learning outperforms Analog SGD significantly and is comparable with Digital SGD.

Table T1.1: Loss (lower is better)

Task 2: Fine-tuning language models. We finetune a pre-training tiny MobilieBERT model with 17M parameters on the MRPC dataset via LoRA training. A total of 462 LoRA layers and 0.99M trainable parameters are introduced. We implemented different portions of the LoRA layers on analog arrays, while the others were left on digital hardware and trained using AdamW. The algorithms reduce to digital Adam if the percentage is 0%, while they reduce to a pure analog algorithm (Analog SGD or Residual Learning) when the percentage is 100%.

Table T1.2: Accuracy (higher is better)

Table T1.3: F1 score (higher is better)

An interesting observation is that Residual Learning achieves slightly higher accuracy than its digital counterpart when the analog percentage is 20%. We speculate that it happens since the analog imperfection implicitly applies a regularization on the training, leading to a better generalization performance. It is worth investigating in future work.

---

Q3. The strong convexity assumption for the main convergence results is quite restrictive for neural networks.

Thank you for pointing it out. Following nonconvex bilevel optimization work [R2.1-R2.2], the strong convexity assumed in the current paper can be relaxed to the Polyak-Lojasiewicz (PL) condition with similar conclusions, which accommodates nonconvex landscapes such as those in overparameterized neural networks [R2.3]. We will discuss it in the revision.

---

Q4. The bounded saturation assumption may not hold as hardware degrades over time.

Thank you for pointing it out. We did not consider hardware degradation over time since typical training processes for moderate-scale models or fine-tuning large models only take hours to days, during which hardware degradation is not a dominant imperfection. It is worth considering in large-scale model training tasks if the training takes a longer time, like weeks to months. We plan to relax this assumption in our future work.

---

Q5. Can the authors validate if the idealized response function model can capture most hardware complexities encountered in practice?

Thanks for your valuable suggestion. As we discuss in Appendix A, a series of empirical works measured the response curves on various resistive elements [45-59], which demonstrate that the non-linear and asymmetric response is one of the main challenges of analog training. On top of that, AIHWKIT [44] provides a fine-grained simulator of the analog hardware, which includes most of the known hardware imperfections like the IO noise, limited granularity, device variation, and quantization error in analog-digital conversion. Building upon their work, the validation in this paper can capture most hardware complexities.

To further validate the proposed idealized response function model, we compare two settings:

• S1) training models following the proposed dynamics (5);
• S2) training models following the standard settings in the AIHWKIT. See Appendix K and Table 4.1 in the response to the Reviewer Uooc for the concrete configurations.

Table T1.4 presents the simulation results. The results show that the accuracy of S1 and S2 is similar for both Analog SGD and Residual Learning. It demonstrates that the proposed response function model captures the training complexities in practice well. In response to your feedback, we will discuss it in the revision.

Table T1.4: Comparison of S1 and S2 in the CIFAR10 training task

---

Q6. The Residual Learning approach introduces computational overhead by requiring maintenance of an additional array and computing gradients on mixed weights.

While Residual Learning introduces extra computational overhead, it is affordable in practice. Compared to Analog SGD, the analog memory requirement doubles, but the latency remains almost unchanged since Residual Learning does not explicitly compute the mixed weights during the forward and backward passes. As [R1.4] suggests, $W_k$ and $P_k$ can share a same analog-digital convertor (ADC), which implements the weight mixing without introducing extra latency. On the other hand, as suggested by [22] in our paper, the forward, backward, and update on $W_k$ and $P_k$ are performed in parallel, which avoids significant latency increase. Consequently, introducing an extra residual array does not incur substantial extra latency. Following the evaluation in Table 1 in [24], we compared the latency of Analog SGD and Residual Learning in Table T1.5. The results suggest that the extra overhead is mild. We will discuss it in the revision.

Table T1.5: Comparison of time (nano-second) consumption in each layer

---

We hope our response resolves your concerns. Thanks!

---

[R1.1] Shen, et al. On penalty-based bilevel gradient descent method. ICML 2023.

[R1.2] Kwon, et al. On penalty methods for nonconvex bilevel optimization and first-order stochastic approximation. ICLR 2024.

[R1.3] Liu, et al. Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. 2022.

[R1.4] Song, et al. Programming memristor arrays with arbitrarily high precision for analog computing. Science 2024.

[R1.5] Jain, et al. A heterogeneous and programmable compute-in-memory accelerator architecture for analog-ai using dense 2-D mesh. VLSI 2022.