解决长序列半马尔可夫模型推理的内存瓶颈问题。
Streaming Structured Inference with Flash-SemiCRF

- 用前缀和数组实时计算边权,避免存储巨大张量。
- 支持流式推理,内存占用随序列长度亚线性增长。
- 适合基因组等超长序列的精确建模,兼容标签不平衡场景。
半马尔可夫条件随机场(semi-CRF)为序列片段分配标签,而非单个位置,可实现片段级特征的精确推断与边界不确定性估计。但现有实现需存储一个随序列长度、最大片段长度和标签数增长的庞大边权张量,在语音级或基因组级(序列超10万位置)问题中变得不可行。本文发现核心瓶颈在于边权张量的显式存储,提出以紧凑前缀和数组替代,并实现三项改进:其一,用前缀和查表取代存储,内存降低倍数与片段长度和标签数乘积成正比;其二,设计带检查点归一化的流式前后向算法,使工作内存亚线性依赖于序列长度,且保持精确梯度;其三,采用零中心累积得分控制数值漂移,并在标签不平衡下诱导自适应持续时间先验。上述方法集成于Flash-SemiCRF——一个融合Triton内核的框架,首次实现此前无法处理的大规模问题的精确半马尔可夫推断。代码已开源:https://github.com/biobenkj/flash-semicrf。
原文摘要 · Abstract (English)
Semi-Markov Conditional Random Fields (semi-CRFs) assign labels to segments of a sequence rather than to individual positions, enabling exact inference over segment-level features and principled uncertainty estimates at their boundaries. However, existing implementations must materialize a large edge potential tensor whose size grows with sequence length, maximum segment length, and label count, becoming prohibitive for speech-scale state spaces and intractable at genomic scales where sequences can exceed 100,000 positions. This memory bottleneck has limited the adoption of exact segment-level inference for long sequences and large label sets. We identify that the core inefficiency is materializing edge potentials that can instead be evaluated on-the-fly from a compact prefix-sum array, and make several improvements. First, replacing the stored edge tensor with prefix-sum lookup reduces the memory footprint by a factor proportional to the product of segment length and label count. Second, a streaming forward-backward pass with checkpoint-boundary normalization keeps working memory sublinear in sequence length while preserving exact gradients. Third, zero-centered cumulative scores control numerical drift and induce an adaptive duration prior under label imbalance. We integrate these ideas into Flash-SemiCRF, a fused Triton kernel that enables exact semi-CRF inference on previously intractable problem sizes. Available at https://github.com/biobenkj/flash-semicrf.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。