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.


翻译:选择性状态空间模型(SSMs)实现了线性时间推理,但其基于梯度的敏感性分析在反向传播过程中仍受限于O(L)的内存扩展,成为性能瓶颈。这一内存约束使得在消费级硬件上进行基因组规模建模(L > 10^5)变得不可行。我们提出了相位梯度流(PGF)框架,该框架通过在状态空间流形上直接运算来计算精确的解析导数,从而避免了构建中间计算图的需求。通过将SSM动力学重新表述为分块算子空间演化(TOSE),我们的方法实现了相对于序列长度的O(1)内存复杂度,与标准自动微分相比,峰值显存使用量减少了94%,吞吐量提高了23倍。与在刚性常微分方程体系中表现出数值发散性的并行前缀扫描不同,PGF通过不变误差缩放确保了稳定性,在极端长序列上仍能保持接近机器精度。我们在包含128,000步序列的脉冲响应基准测试中验证了PGF的实用性——在此规模下,传统的自动微分会遇到极高的内存开销,并通常在多层模型中导致内存不足(OOM)故障。我们的工作使得在单块GPU上进行染色体规模的敏感性分析成为可能,从而弥合了理论上的无限上下文模型与实际硬件限制之间的鸿沟。

0
下载
关闭预览

相关内容

【NeurIPS2019】图变换网络:Graph Transformer Network
CosFace: Large Margin Cosine Loss for Deep Face Recognition论文笔记
统计学习与视觉计算组
44+阅读 · 2018年4月25日
国家自然科学基金
0+阅读 · 2017年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
VIP会员
相关基金
国家自然科学基金
0+阅读 · 2017年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员