[论文解读] Mixed-Precision Training and Compilation for RRAM-based Computing-in-Memory Accelerators

Rebecca Pelke, Joel Klein+4 more|arXiv (Cornell University)|Jan 29, 2026

Parallel Computing and Optimization Techniques被引用 0

一句话总结

引入了一种混合精度量化与强化学习驱动的编译器框架（CIM-AQ）用于RRAM CIM加速器，在最小精度损失的情况下实现显著的低延迟，同时提供基于TVM、Brevitas作为QAT后端的CIM支持的MPQ编译器。

ABSTRACT

Computing-in-Memory (CIM) accelerators are a promising solution for accelerating Machine Learning (ML) workloads, as they perform Matrix-Vector Multiplications (MVMs) on crossbar arrays directly in memory. Although the bit widths of the crossbar inputs and cells are very limited, most CIM compilers do not support quantization below 8 bit. As a result, a single MVM requires many compute cycles, and weights cannot be efficiently stored in a single crossbar cell. To address this problem, we propose a mixed-precision training and compilation framework for CIM architectures. The biggest challenge is the massive search space, that makes it difficult to find good quantization parameters. This is why we introduce a reinforcement learning-based strategy to find suitable quantization configurations that balance latency and accuracy. In the best case, our approach achieves up to a 2.48x speedup over existing state-of-the-art solutions, with an accuracy loss of only 0.086 %.

研究动机与目标

• 鉴于交叉开关单元位宽和DAC/ADC约束，说明对RRAM CIM加速器进行低精度量化的必要性。
• 提出面向CIM架构的混合精度量化与训练框架（MPQ），以降低延迟。
• 开发一个基于强化学习的CIM感知自动量化优化器（CIM-AQ），以高效探索MPQ搜索空间。
• 在基于TVM的流水线中创建一个支持Brevitas训练的MPQ模型映射至CIM目标的编译器。
• 在ImageNet上对标准的CNN/视觉模型演示在受控精度损失情况下的延迟改进。

提出的方法

• 将Hardware-aware Automated Quantization（HAQ）框架扩展到CIM目标，形成CIM-AQ，在CIM硬件上优化延迟-精度权衡。
• 定义一个新的奖励函数，目标是最小精度acc_t = acc_8b - acc_loss，同时对低于目标精度进行惩罚并优先考虑延迟降低。
• 使用交叉开关参数（M×N、单元精度、DAC精度、写入和MVM时间）来建模CIM延迟，并推导逐层和逐配置的延迟估算。
• 在QAT后端使用Brevitas实现MPQ，以支持更广泛的网络层（包括 transformers）并加速基于RL的优化。
• 在TVM上开发面向CIM的编译器：将QAT训练模型转换为ONNX，应用QDQ风格量化，执行图转换（QNNFuse），并在CPU与CIM之间分配工作负载，使用自定义降低和分阶段缓冲区。
• 引入权重驻留数据流和逐层轴标 Scheduling，以优化交叉开关利用率并最小化写入。

实验结果

研究问题

• RQ1 如何在CIM目标上有效学习混合精度量化以在保持目标精度的同时最小化延迟？
• RQ2 在指导MPQ参数搜索（如激活/权重量化位宽、单元分辨率）以实现延迟优化推理方面，哪些CIM特有约束或策略最有效？
• RQ3 统一的CIM感知强化学习优化器（CIM-AQ）是否在速度和准确性方面超过固定8位CIM量化？
• RQ4 如何将Brevitas整合到CIM专注的RL优化循环中，以扩大网络支持并提升运行时性能？
• RQ5 基于TVM的CIM编译器在将MPQ模型映射到RRAM CIM硬件用于CNN和Transformer方面的有效性如何？
• RQ6 不同的交叉开关配置（单元分辨率、交叉开关大小）在MPQ下对延迟-精度权衡的影响如何？

主要发现

• 该框架在8位CIM基线下实现最高2.48×的加速，精度损失为0.086%。
• CIM-AQ使多种架构（包括Transformer）的混合精度优化成为可能，在RL训练中相较于之前的HAQ优化在GPU上大约快8%。
• 激活位宽在CIM中往往比权重更积极量化，这是由延迟考虑驱动的。
• 约束MPQ搜索空间（输入/输出、权重或两者约束）可在S/AL分数上改善速度/精度权衡，在ResNet-18上组合约束获得最佳结果。
• 对于ViT-B/32，MPQ方法在各层均实现显著加速且精度损失在可接受范围内，在所有测试模型中观察到至少2.20×的加速。