🤖 machine learning

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

ThRIve is a noise-aware training methodology that leverages low-rank adaptation to store critical parameters on thermally stable hardware, enabling CNN inference on heterogeneous PIM architectures to maintain high accuracy across varying temperatures while achieving up to a 5.4x reduction in energy-delay product compared to SRAM-based systems.

Original authors: Vibhanshu Sharma, Pratyush Dhingra, Janardhan Rao Doppa, Partha Pratim Pande

Published 2026-07-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Vibhanshu Sharma, Pratyush Dhingra, Janardhan Rao Doppa, Partha Pratim Pande

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build a super-fast, super-efficient brain for your phone or a smart robot. For a long time, computers have been like a librarian who has to run back and forth between the library shelves (memory) and the reading desk (processor) to get books. This running back and forth takes a lot of energy and time. Scientists have been dreaming of a "Processing-in-Memory" (PIM) system, where the librarian sits right on the shelf and does the reading and thinking without ever leaving. This is incredibly efficient, especially for the heavy math needed to run Artificial Intelligence (AI).

However, there's a catch. The new kind of "shelves" scientists are using to build these brains are made of special materials that are great at storing data but are very sensitive to heat. Think of them like chocolate bars left in the sun; as they get warmer, the chocolate starts to melt and lose its shape. In these computer chips, heat causes the stored information (the AI's "knowledge") to wiggle and distort. This "thermal noise" makes the AI start making mistakes, like a student who knows the answer but gets distracted by a buzzing fly. The big question is: How do we keep these super-efficient AI brains working perfectly, even when they get hot and the information starts to melt?

This is where a team of researchers steps in with a clever solution called ThRIve. They realized that trying to fix the melting chocolate by cooling the whole room is too slow and expensive. Instead, they came up with a two-part strategy that mixes hardware and software. First, they use a technique called Low-Rank Adaptation (LoRA). Imagine the AI model is a giant, heavy encyclopedia. Instead of rewriting the whole book every time the temperature changes (which would be slow and wear out the paper), they write tiny, flexible sticky notes with corrections on them. These sticky notes are small, easy to move, and can be updated quickly.

The second part of the trick is where they put these sticky notes. The heavy encyclopedia stays on the "melting" shelves (the heat-sensitive memory), but the tiny, crucial sticky notes are stored on a different, super-sturdy shelf made of a material that doesn't melt in the heat (a type of memory called SRAM). The AI reads the heavy book and the sticky notes at the same time, combining them to get the right answer. To make sure these sticky notes work for any temperature, the researchers used a smart search method called Bayesian Optimization. Instead of guessing random temperatures to train the AI, this method acts like a detective, hunting down the specific "worst-case" heat scenarios and teaching the AI how to handle them efficiently.

The results of their experiments are quite promising. When they tested this ThRIve system on different AI models and hardware setups, they found that the AI could maintain its accuracy even as the temperature swung wildly. In fact, the average accuracy stayed within 2% of what it would be in a perfect, noise-free world, and the variation in accuracy across the whole temperature range stayed within 2% of that average. This means the AI didn't get confused or forget things just because it got hot. Furthermore, by keeping the heavy data on the efficient, heat-sensitive shelves and only using the sturdy shelves for the tiny corrections, they saved a massive amount of energy and time. In some tests, the system was up to 5.4 times more efficient (in terms of energy-delay product) than using only the sturdy, non-melting shelves.

The researchers also showed that older methods of trying to fix this problem—like trying to retrain the whole AI model for every single temperature or just hoping the AI could adapt on the fly—didn't work as well. Those older approaches often led to the AI "forgetting" how to handle cooler temperatures while trying to learn hot ones. ThRIve, by contrast, seems to generalize well, handling the entire range of temperatures without needing to be constantly retrained. While these findings come from simulations and computer models rather than physical chips in a real-world device, the data suggests that this mix of smart software and mixed-material hardware could be a key to making reliable, energy-efficient AI for our future devices.

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

1. Problem Statement

Processing-in-Memory (PIM) architectures, particularly those utilizing Non-Volatile Memory (NVM) devices like ReRAM, PCM, and FeFET, offer significant energy efficiency and performance advantages for accelerating Machine Learning (ML) workloads by performing Matrix-Vector Multiplication (MVM) directly within memory arrays. However, these NVM-based devices suffer from inherent non-idealities, with thermal noise being a critical challenge. Thermal noise alters the stored values in memory cells (which correspond to model weights) as temperature fluctuates, leading to a degradation in inference accuracy.

Existing mitigation strategies, such as Noise-Aware Training (NAT), attempt to adapt model weights to specific noise levels. However, these methods face two primary limitations:

Lack of Generalizability: Current NAT techniques are typically trained offline on GPUs at discrete temperature points (e.g., 25°C, 35°C). When deployed on NVM crossbars, these "noise-aware" weights are still exposed to thermal noise at unseen temperatures, causing further deviation and accuracy loss.
Hardware Overhead and Endurance: Storing all model parameters in a thermally resilient environment (like SRAM) to ensure robustness incurs prohibitive hardware overhead. Conversely, frequent on-device retraining to adapt to changing temperatures is impractical due to the limited write endurance of NVM devices.

2. Methodology: ThRIve Framework

The authors propose ThRIve, a hardware-software co-design methodology that leverages Low-Rank Adaptation (LoRA) within heterogeneous NVM-CMOS PIM architectures to achieve thermally robust inference.

2.1 Hardware Strategy: Heterogeneous Storage

ThRIve exploits the complementary strengths of different memory technologies:

NVM (ReRAM, PCM, FeFET): Stores the original, pre-trained model base weights. These provide high density and energy efficiency but are susceptible to thermal noise.
SRAM (CMOS-based): Stores the noise-aware LoRA parameters. SRAM is thermally resilient and less susceptible to temperature-induced weight deviations.

During inference, the system computes the output as the sum of two parallel paths:

The base model computation using noisy weights on the NVM crossbar.
The LoRA adaptation computation using stable weights on the SRAM crossbar.
The final activation is derived from the summation of these outputs, effectively correcting the thermal noise-induced errors without altering the base NVM weights.

2.2 Algorithmic Strategy: Bayesian Optimization-Driven NAT

To train the model effectively across a continuous temperature range without exhaustive retraining, ThRIve employs a min-max optimization framework driven by Bayesian Optimization (BO):

LoRA Parameterization: The model is modified to include trainable low-rank matrices ( $A$ and $B$ ) added to the frozen base weights. The effective weight becomes $W_{effective} = W_{noisy} + BA$ .
Min-Max Optimization: The goal is to find LoRA parameters that minimize prediction loss even under the "worst-case" thermal noise conditions.
- Maximization (Inner Loop): BO adaptively searches the continuous temperature/noise space to identify the specific noise level (temperature) that causes the maximum prediction loss for the current model.
- Minimization (Outer Loop): The model is re-trained using LoRA backpropagation specifically on this identified worst-case noise level to reduce the loss.
Efficiency: This iterative process avoids the need to train on every possible temperature point or rely on sequential learning, which often suffers from catastrophic forgetting.

3. Key Contributions

The paper outlines three primary contributions:

LoRA Integration for Thermal Robustness: The proposal to use LoRA within heterogeneous NVM-CMOS architectures. This significantly reduces the storage overhead required for thermal robustness, as only a small fraction of parameters (the low-rank matrices) need to be stored on the thermally resilient SRAM, while the bulk of the model remains on energy-efficient NVM.
Efficient BO-Driven NAT: The development of a Bayesian Optimization-driven training methodology. This approach adaptively selects representative noise levels for training, allowing the LoRA matrices to generalize to unseen temperatures without altering base parameters or requiring frequent on-device retraining, thereby preserving NVM write endurance.
Performance and Robustness Validation: Experimental demonstration that ThRIve-enabled architectures maintain consistent inference accuracy across the entire operating temperature range (300K–400K), outperforming existing NAT methods and achieving performance comparable to all-SRAM systems but with significantly lower energy costs.

4. Experimental Results

The methodology was evaluated using CNN models (VGG16, ResNet18, MobileNetv3_Small) on datasets (CIFAR-10, CIFAR-100, TinyImageNet) across heterogeneous PIM architectures (HyperX, HyDE, HuNT).

Accuracy: ThRIve-enabled architectures maintained a mean inference accuracy within 2% of the ideal (noise-free) accuracy. Furthermore, the variation in accuracy across the entire operating temperature range remained within 2% of the mean value. This consistency was achieved across ReRAM, PCM, and FeFET-based systems.
Comparison to Baselines: Unlike existing NAT methods (e.g., TRRIMC, BN OTF) which showed significant accuracy drops and high range deviation due to catastrophic forgetting and lack of generalization, ThRIve provided stable performance across all temperatures.
Energy-Delay Product (EDP): ThRIve achieved accuracy and robustness comparable to thermally-resilient SRAM-based PIM systems but delivered substantial efficiency gains:
- Up to ~5.4× reduction in EDP on HuNT architectures.
- Up to ~4.8× reduction in EDP on HyperX architectures.
- Up to ~2.8× reduction in EDP on HyDE architectures.
Rank Analysis: The study determined optimal rank proportions ( $P$ ) for LoRA matrices (e.g., $P=1/12$ for ReRAM, $1/14$ for PCM, $1/16$ for FeFET) to balance accuracy and EDP overhead.

5. Significance and Claims

The paper claims that ThRIve fills a critical gap in the state-of-the-art by enabling generalized thermal robustness for NVM-based PIM systems. By decoupling the storage of noise-aware parameters from the noisy memory substrate, ThRIve eliminates the need for frequent on-device retraining and the associated endurance penalties.

The authors assert that their approach makes heterogeneous PIM architectures viable for real-world edge AI deployment, where operating temperatures fluctuate. The methodology offers a practical trade-off, delivering the thermal resilience of all-SRAM systems while retaining the high density and energy efficiency of NVM technologies, ultimately enabling reliable ML inference in thermally variable environments without compromising on energy efficiency.