Inferência Estruturada em Streaming com Flash-SemiCRF

Streaming Structured Inference with Flash-SemiCRF

April 20, 2026

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

cs.AI

Resumo

Os Campos Condicionais Aleatórios Semi-Markov (semi-CRFs) atribuem rótulos a segmentos de uma sequência em vez de a posições individuais, permitindo inferência exata sobre características a nível de segmento e estimativas de incerteza fundamentadas nos seus limites. No entanto, as implementações existentes precisam materializar um grande tensor de potencial de aresta, cujo tamanho aumenta com o comprimento da sequência, o comprimento máximo do segmento e a quantidade de rótulos, tornando-se proibitivo para espaços de estado de escala de fala e intratável em escalas genómicas, onde as sequências podem exceder 100.000 posições. Este estrangulamento de memória tem limitado a adoção da inferência exata a nível de segmento para sequências longas e grandes conjuntos de rótulos. Nós identificamos que a ineficiência central é a materialização de potenciais de aresta que podem, em vez disso, ser avaliados dinamicamente a partir de um *array* compacto de soma de prefixos, e realizamos várias melhorias. Primeiro, substituir o tensor de aresta armazenado pela consulta à soma de prefixos reduz a pegada de memória por um fator proporcional ao produto do comprimento do segmento e da quantidade de rótulos. Segundo, um passo *forward-backward* em fluxo contínuo com normalização nos limites dos pontos de verificação mantém a memória de trabalho sublinear em relação ao comprimento da sequência, preservando os gradientes exatos. Terceiro, pontuações cumulativas com centro em zero controlam o desvio numérico e induzem um prévio de duração adaptativo sob desequilíbrio de rótulos. Integramos estas ideias no Flash-SemiCRF, um núcleo (*kernel*) Triton fundido que permite a inferência exata de semi-CRFs em tamanhos de problema anteriormente intratáveis. Disponível em 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.