Streaming Gestructureerde Inferentie met Flash-SemiCRF
Streaming Structured Inference with Flash-SemiCRF
April 20, 2026
Auteurs: Benjamin K. Johnson, Thomas Goralski, Ayush Semwal, Hui Shen, H. Josh Jang
cs.AI
Samenvatting
Semi-Markov Conditionele Random Fields (semi-CRFs) kennen labels toe aan segmenten van een sequentie in plaats van aan individuele posities, wat exacte inferentie over segment-level kenmerken en geprincipeerde onzekerheidsschattingen aan hun grenzen mogelijk maakt. Bestaande implementaties moeten echter een grote edge potential tensor materialiseren, waarvan de grootte toeneemt met de sequentielengte, de maximale segmentlengte en het aantal labels, wat onhaalbaar wordt voor state spaces op spraakschaal en onhanteerbaar op genomische schaal, waar sequenties meer dan 100.000 posities kunnen overschrijden. Deze geheugenflessenhals heeft de adoptie van exacte segment-level inferentie voor lange sequenties en grote labelsets beperkt. Wij identificeren dat de kerninefficiëntie ligt in het materialiseren van edge potentials die in plaats daarvan on-the-fly kunnen worden geëvalueerd vanuit een compact prefix-sum array, en doen verschillende verbeteringen. Ten eerste reduceert het vervangen van de opgeslagen edge tensor door prefix-sum lookup de geheugenvoetafdruk met een factor evenredig aan het product van segmentlengte en labelaantal. Ten tweede houdt een streaming forward-backward pass met checkpoint-boundary normalisatie het werkgeheugen sublineair in de sequentielengte, terwijl exacte gradienten behouden blijven. Ten derde controleren zero-centered cumulative scores numerieke drift en induceren ze een adaptieve duration prior onder labelonbalans. Wij integreren deze ideeën in Flash-SemiCRF, een gefuseerde Triton-kernel die exacte semi-CRF inferentie op voorheen onhanteerbare probleemgroottes mogelijk maakt. Beschikbaar op 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.