State Space Models 更适合在设备端处理超长文本,性能远超 Transformer。
Characterizing State Space Model and Hybrid Language Model Performance with Long Context
- 采用 SSM 架构,计算复杂度线性增长,支持百万级标记输入。
- 在约 5.7 万标记时,SSM 比 Transformer 快达 4 倍,内存占用减少 64%。
- 自定义 SSM 算子(如 selective scan)成主要延迟来源,占推理时间超 55%。
增强现实等新兴应用正推动本地设备上对连续或长上下文输入的机器智能需求。然而,当前主流的 Transformer 模型存在二次方级计算与内存开销,难以支撑长文本处理。这促使研究转向 State Space Models (SSMs) 及 SSM-Transformer 混合模型,其具备近线性扩展性,可高效处理数百万标记并保持高性能。尽管已有研究展示前景,但这些模型在计算性能与硬件资源需求方面的负载特性尚未充分探索,限制了系统级优化的理解。为此,本文针对消费级与嵌入式 GPU 上的长上下文推理,对精选的 Transformer、SSM 与混合模型进行了全面对比基准测试。结果表明,SSM 在消费级与嵌入式 GPU 上更适合作为长上下文推理的本地 AI 架构。虽然在短序列(<8K 标记)下,Transformer 最快可达 1.9 倍加速;但在极长上下文(~57K 标记)下,SSM 性能出现显著反转,最快达 4 倍加速,且内存占用降低约 64%。算子级分析显示,尽管 selective scan 等自定义 SSM 内核已针对硬件优化以减少内存读写,但其串行、逐元素特性仍主导推理延迟,占边缘平台运行时间超 55%。代码已开源:https://github.com/sapmitra/ssm-scope。
原文摘要 · Abstract (English)
Emerging applications such as AR are driving demands for machine intelligence capable of processing continuous and/or long-context inputs on local devices. However, currently dominant models based on Transformer architecture suffers from the quadratic computational and memory overhead, which hinders applications required to process long contexts. This has spurred a paradigm shift towards new architectures like State Space Models (SSMs) and SSM-Transformer hybrid models, which provide near-linear scaling. The near-linear scaling enabled efficient handling of millions of tokens while delivering high performance in recent studies. Although such works present promising results, their workload characteristics in terms of computational performance and hardware resource requirements are not yet thoroughly explored, which limits our understanding of their implications to the system level optimizations. To address this gap, we present a comprehensive, compara-ive benchmarking of carefully selected Transformers, SSMs, and hybrid models specifically for long-context inference on consumer and embedded GPUs. Our analysis shows that SSMs are well-suited for on-device AI on consumer and embedded GPUs for long context inferences. While Transformers are up to 1.9x faster at short sequences (<8K tokens), SSMs demonstrate a dramatic performance inversion, becoming up to 4x faster at very long contexts (~57K tokens), thanks to their linear computational complexity and ~64% reduced memory footrprint. Our operator-level analysis reveals that custom SSM kernels like selective scan despite being hardware-aware to minimize memory IO, dominate the inference runtime on edge platforms, accounting for over 55% of latency due to their sequential, element-wise nature. SSM-Scope is open-sourced at https://github.com/sapmitra/ssm-scope
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。