Inferencia Estructurada en Streaming con 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
Resumen
Los Campos Aleatorios Condicionales Semi-Markov (semi-CRFs) asignan etiquetas a segmentos de una secuencia en lugar de a posiciones individuales, lo que permite realizar inferencia exacta sobre características a nivel de segmento y estimaciones de incertidumbre fundamentadas en sus límites. Sin embargo, las implementaciones existentes deben materializar un gran tensor de potenciales de arista cuyo tamaño aumenta con la longitud de la secuencia, la longitud máxima del segmento y el número de etiquetas, volviéndose prohibitivo para espacios de estados a escala de habla e intratable a escalas genómicas donde las secuencias pueden superar las 100.000 posiciones. Este cuello de botella de memoria ha limitado la adopción de la inferencia exacta a nivel de segmento para secuencias largas y grandes conjuntos de etiquetas. Identificamos que la ineficiencia central es la materialización de potenciales de arista que, en su lugar, pueden evaluarse sobre la marcha a partir de un array compacto de suma de prefijos, y realizamos varias mejoras. En primer lugar, reemplazar el tensor de aristas almacenado con la búsqueda por suma de prefijos reduce la huella de memoria en un factor proporcional al producto de la longitud del segmento y el número de etiquetas. En segundo lugar, un pase hacia adelante y hacia atrás en flujo continuo con normalización en los límites de verificación mantiene la memoria de trabajo sublineal con respecto a la longitud de la secuencia, preservando al mismo tiempo los gradientes exactos. En tercer lugar, las puntuaciones acumuladas centradas en cero controlan la deriva numérica e inducen un prior de duración adaptativo bajo desequilibrio de etiquetas. Integramos estas ideas en Flash-SemiCRF, un kernel de Triton fusionado que permite la inferencia exacta de semi-CRF en tamaños de problema previamente intratables. Disponible en 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.