ThRIve: Thermally Robust CNN Inference via Low-Rank Adaptation in Heterogeneous PIM Architectures

Processing-In-Memory (PIM) architectures using Non-Volatile Memory (NVM) are highly efficient for machine learning, but they are inherently sensitive to thermal noise. As temperatures fluctuate, the physical properties of NVM cells change, causing weight deviations that degrade neural network inference accuracy. While existing noise-aware training (NAT) methods attempt to mitigate this, they often require frequent, energy-intensive retraining or lack generalizability across the full operating temperature range. Furthermore, storing all model parameters in thermally resilient memory (like SRAM) to avoid these issues is prohibitively expensive in terms of hardware area.

ThRIve addresses these challenges by combining Low-Rank Adaptation (LoRA) with a Bayesian Optimization-driven training framework. Instead of retraining the entire model, ThRIve freezes the base model weights and trains small, low-rank matrices that adapt to temperature-induced noise. During deployment, these compact LoRA parameters are stored in CMOS-based memory (which is naturally resistant to thermal noise), while the bulk of the model remains in the high-density NVM.

To optimize the training process, the authors utilize Bayesian Optimization to select a representative subset of temperature-dependent noise levels. This ensures the model learns to generalize across the entire operating range without requiring exhaustive training at every possible temperature point. This approach significantly reduces the need for on-device retraining, thereby extending the lifespan of NVM cells by minimizing write operations.

Experimental results show that ThRIve maintains inference accuracy within 2% of ideal, noise-free performance across the entire operating temperature range. By offloading only the LoRA parameters to resilient hardware, the system achieves robustness comparable to pure SRAM-based PIM systems but with significantly higher storage density. Additionally, the framework delivers up to a 5.4x reduction in the energy-delay product (EDP) compared to traditional approaches, making it a highly efficient solution for deploying reliable AI at the edge.