Inferenza Strutturata in Streaming con Flash-SemiCRF

Streaming Structured Inference with Flash-SemiCRF

April 20, 2026

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

cs.AI

Abstract

I Campi Casuali Condizionali Semi-Markoviani (semi-CRFs) assegnano etichette a segmenti di una sequenza piuttosto che a singole posizioni, consentendo inferenza esatta su feature a livello di segmento e stime di incertezza rigorose ai loro confini. Tuttavia, le implementazioni esistenti devono materializzare un ampio tensore di potenziali sugli archi la cui dimensione cresce con la lunghezza della sequenza, la lunghezza massima del segmento e il numero di etichette, diventando proibitiva per spazi degli stati di scala vocale e intrattabile a scale genomiche dove le sequenze possono superare le 100.000 posizioni. Questo collo di bottiglia della memoria ha limitato l'adozione dell'inferenza esatta a livello di segmento per sequenze lunghe e grandi insiemi di etichette. Identifichiamo che l'inefficienza centrale è la materializzazione di potenziali sugli archi che possono invece essere valutati al volo da un array compatto di somme prefisse, e apportiamo diversi miglioramenti. In primo luogo, sostituire il tensore degli archi memorizzato con la ricerca per somme prefisse riduce l'ingombro di memoria di un fattore proporzionale al prodotto della lunghezza del segmento e del numero di etichette. In secondo luogo, un passaggio in avanti e all'indietro in streaming con normalizzazione ai checkpoint mantiene la memoria di lavoro sublineare rispetto alla lunghezza della sequenza preservando gradienti esatti. In terzo, punteggi cumulativi a media zero controllano la deriva numerica e inducono un prior adattivo sulla durata in caso di squilibrio delle etichette. Integriamo queste idee in Flash-SemiCRF, un kernel Triton fuso che abilita l'inferenza esatta per semi-CRF su dimensioni di problema precedentemente intrattabili. Disponibile all'indirizzo 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.