Потоковый структурированный вывод с Flash-SemiCRF

Streaming Structured Inference with Flash-SemiCRF

April 20, 2026

Авторы: Benjamin K. Johnson, Thomas Goralski, Ayush Semwal, Hui Shen, H. Josh Jang

cs.AI

Аннотация

Полу-марковские условные случайные поля (semi-CRFs) присваивают метки сегментам последовательности, а не отдельным позициям, что позволяет проводить точный вывод на основе признаков на уровне сегментов и получать обоснованные оценки неопределенности на их границах. Однако существующие реализации требуют материализации большого тензора потенциалов ребер, размер которого растет с длиной последовательности, максимальной длиной сегмента и количеством меток, что становится непозволительно затратным для пространств состояний речевого масштаба и неразрешимо на геномных масштабах, где последовательности могут превышать 100 000 позиций. Это узкое место по памяти ограничивало применение точного вывода на уровне сегментов для длинных последовательностей и больших наборов меток. Мы выявили, что основная неэффективность заключается в материализации потенциалов ребер, которые вместо этого можно вычислять на лету из компактного массива префиксных сумм, и внесли несколько улучшений. Во-первых, замена хранимого тензора ребер на поиск по префиксным суммам сокращает объем памяти пропорционально произведению длины сегмента и количества меток. Во-вторых, потоковый прямой-обратный проход с нормализацией по контрольным точкам-границам поддерживает объем рабочей памяти сублинейным относительно длины последовательности, сохраняя при этом точные градиенты. В-третьих, центрированные относительно нуля кумулятивные scores контролируют численный дрейф и индуцируют адаптивный априор длительности при дисбалансе меток. Мы интегрируем эти идеи в Flash-SemiCRF — объединенное ядро Triton, которое позволяет проводить точный вывод semi-CRF для ранее неразрешимых размеров задач. Доступно по адресу 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.