Flash-SemiCRF를 활용한 스트리밍 구조 추론

초록

준-마르코프 조건부 무작위 필드(semi-CRFs)는 개별 위치가 아닌 시퀀스의 세그먼트에 레이블을 할당하여 세그먼트 수준 특징에 대한 정확한 추론과 경계에서의 이론적으로 타당한 불확실성 추정을 가능하게 합니다. 그러나 기존 구현은 시퀀스 길이, 최대 세그먼트 길이, 레이블 수에 따라 크기가 증가하는 대규모 에지 포텐셜 텐서를 구체화해야 하며, 이는 음성 규모의 상태 공간에서는 실용적이지 않고 서열 길이가 100,000개 위치를 초과할 수 있는 게놈 규모에서는 다루기 어렵습니다. 이러한 메모리 병목 현상으로 인해 긴 시퀀스와 큰 레이블 집합에 대한 정확한 세그먼트 수준 추론의 채택이 제한되어 왔습니다. 우리는 핵심 비효율성이 구체화되는 에지 포텐셜에 있으며, 이는 컴팩트한 접두사 합 배열을 통해 실시간으로 평가될 수 있음을 확인하고 몇 가지 개선을 이루었습니다. 첫째, 저장된 에지 텐서를 접두사 합 조회로 대체하면 메모리 사용량이 세그먼트 길이와 레이블 수의 곱에 비례하여 감소합니다. 둘째, 체크포인트-경계 정규화를 사용한 스트리밍 순방향-역방향 전파는 정확한 기울기를 유지하면서 작업 메모리를 시퀀스 길이에 대해 하위 선형으로 유지합니다. 셋째, 영점 중심 누적 점수는 수치적 드리프트를 제어하고 레이블 불균형 하에서 적응형 지속 시간 사전 분포를 유도합니다. 우리는 이러한 아이디어를 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.

Flash-SemiCRF를 활용한 스트리밍 구조 추론

Streaming Structured Inference with Flash-SemiCRF

초록

Support