提出PGF框架，让状态空间模型实现零内存增长的精确梯度计算。
Breaking the Memory Wall: Exact Analytical Differentiation via Tiled Operator-Space Evolution
- 在状态空间流形上直接计算梯度，不保存中间计算图
- 序列长度128,000时，显存降低94%，吞吐提升23倍
- 适合基因组级长序列建模，单卡即可运行染色体尺度分析
选择性状态空间模型（SSMs）可实现线性时间推理，但其基于梯度的敏感性分析受反向传播中O(L)内存开销限制，难以在消费级硬件上进行超长序列（L > 10^5）建模。本文提出相位梯度流（PGF），通过在状态空间流形上直接计算精确解析导数，避免中间计算图的显式存储。将SSM动态重构成分块算子空间演化（TOSE），实现相对于序列长度的O(1)内存复杂度，在128,000步脉冲响应基准测试中，相比标准Autograd减少94%峰值显存占用，吞吐提升23倍。与存在数值发散问题的并行前缀扫描不同，PGF通过不变误差缩放保证稳定性，维持近机器精度。本方法使单个GPU即可完成染色体尺度敏感性分析，弥合理论无限上下文模型与实际硬件约束之间的鸿沟。
原文摘要 · Abstract (English)
Selective State Space Models (SSMs) achieve linear-time inference, yet their gradient-based sensitivity analysis remains bottlenecked by O(L) memory scaling during backpropagation. This memory constraint precludes genomic-scale modeling (L > 10^5) on consumer-grade hardware. We introduce Phase Gradient Flow (PGF), a framework that computes exact analytical derivatives by operating directly in the state-space manifold, bypassing the need to materialize the intermediate computational graph. By reframing SSM dynamics as Tiled Operator-Space Evolution (TOSE), our method delivers O(1) memory complexity relative to sequence length, yielding a 94% reduction in peak VRAM and a 23x increase in throughput compared to standard Autograd. Unlike parallel prefix scans that exhibit numerical divergence in stiff ODE regimes, PGF ensures stability through invariant error scaling, maintaining near-machine precision across extreme sequences. We demonstrate the utility of PGF on an impulse-response benchmark with 128,000-step sequences - a scale where conventional Autograd encounters prohibitive memory overhead, often leading to out-of-memory (OOM) failures in multi-layered models. Our work enables chromosome-scale sensitivity analysis on a single GPU, bridging the gap between theoretical infinite-context models and practical hardware limitations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。