Inférence Structurée en Streaming avec 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
Résumé
Les champs conditionnels aléatoires semi-markoviens (semi-CRF) attribuent des étiquettes à des segments d'une séquence plutôt qu'à des positions individuelles, permettant une inférence exacte sur les caractéristiques au niveau des segments et des estimations probabilistes rigoureuses de leurs limites. Cependant, les implémentations existantes doivent matérialiser un grand tenseur de potentiels d'arêtes dont la taille augmente avec la longueur de la séquence, la longueur maximale des segments et le nombre d'étiquettes, devenant prohibitif pour les espaces d'états à l'échelle de la parole et insoluble aux échelles génomiques où les séquences peuvent dépasser 100 000 positions. Ce goulot d'étranglement mémoire a limité l'adoption de l'inférence exacte au niveau des segments pour les longues séquences et les grands ensembles d'étiquettes. Nous identifions que l'inefficacité principale est la matérialisation des potentiels d'arêtes qui peuvent plutôt être évalués à la volée à partir d'un tableau compact de sommes préfixées, et apportons plusieurs améliorations. Premièrement, le remplacement du tenseur d'arêtes stocké par une recherche de sommes préfixées réduit l'empreinte mémoire d'un facteur proportionnel au produit de la longueur des segments et du nombre d'étiquettes. Deuxièmement, une passe avant-arrière en flux continu avec normalisation aux limites de vérification maintient la mémoire de travail sous-linéaire par rapport à la longueur de la séquence tout en préservant les gradients exacts. Troisièmement, des scores cumulatifs centrés sur zéro contrôlent la dérive numérique et induisent un a priori de durée adaptatif sous déséquilibre des étiquettes. Nous intégrons ces idées dans Flash-SemiCRF, un noyau Triton fusionné qui permet l'inférence exacte de semi-CRF sur des tailles de problèmes précédemment insolubles. Disponible à l'adresse 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.