フラッシュ半条件付き確率場によるストリーミング構造化推論

Streaming Structured Inference with Flash-SemiCRF

April 20, 2026

著者: Benjamin K. Johnson, Thomas Goralski, Ayush Semwal, Hui Shen, H. Josh Jang

cs.AI

要旨

半マルコフ条件付き確率場（semi-CRFs）は、系列の個々の位置ではなくセグメント単位でラベルを割り当てることで、セグメントレベル特徴量に対する厳密な推論と境界における理論的に妥当な不確実性推定を可能にする。しかし、既存の実装では系列長・最大セグメント長・ラベル数に比例して大きくなる大規模なエッジポテンシャルテンソルを具現化する必要があり、音声処理レベルの状態空間では計算が困難となり、10万位置を超えるゲノム規模系列では処理不能となる。このメモリボトルネックが、長系列と大規模ラベル集合に対する厳密なセグメントレベル推論の普及を妨げてきた。本研究では、本質的な非効率性が、コンパクトな接頭辞和配列からオンザフライ評価可能なエッジポテンシャルの具現化にあることを特定し、三つの改善を提案する。第一に、保存済みエッジテンソルを接頭辞和検索に置換することで、メモリ使用量をセグメント長とラベル数の積に比例して削減する。第二に、チェックポイント境界正規化を伴うストリーミング前後方向アルゴリズムにより、作業メモリを系列長に対して劣線形に保ちつつ厳密な勾配を維持する。第三に、ゼロ中心化累積スコアが数値的ドリフトを抑制し、ラベル不均衡下で適応的期間事前分布を誘導する。これらのアイデアを統合したFlash-SemiCRFは、融合型Tritonカーネルとして実装され、従来扱えなかった問題規模での厳密なsemi-CRF推論を実現する。https://github.com/biobenkj/flash-semicrf で公開。

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.