IMSSA:在 memristive in-memory 计算硬件上部署现代状态空间模型
机器学习
2025-07-09 v1 硬件体系结构
摘要
处理长时序序列是深度学习中的关键挑战。近年来,Transformer 成为该任务的先进方法,但因需显式存储序列而导致内存需求过大。为此,近期涌现出结构化状态空间顺序模型(S4),在保持固定内存状态的同时仍能处理极长序列上下文。S4 模型中状态的循环线性更新使其在现代图形处理器(GPU)上高度高效,通过将循环展开为卷积实现。然而,这一方法需要大量内存和大规模并行计算,仅能在最新的 GPU 上实现。本文旨在将 S4 模型的优势带到边缘硬件,通过对 S4D 模型进行量化感知训练显著降低其规模和计算需求,甚至实现简单实际任务的三值权重。为此,我们扩展常规量化感知训练,使其针对模拟 in-memory 计算硬件优化。随后,我们在 memristive 交叉栏阵列上实现循环 S4D 核,实现其 in-memory 计算式计算。据我们了解,这是首个在 in-memory 计算硬件上实现 S4 核的实现。
引用
@article{arxiv.2412.20215,
title = {IMSSA: Deploying modern state-space models on memristive in-memory compute hardware},
author = {Sebastian Siegel and Ming-Jay Yang and John-Paul Strachan},
journal= {arXiv preprint arXiv:2412.20215},
year = {2025}
}
备注
5 pages, 4 figures, submitted to IEEE ISCAS 2025