PaperlayerView exact revision on arXiv

Source-linked AI summary

Streaming Structured Inference with Flash-SemiCRF

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

arXiv:2604.18780v1cs.LG

TL;DR

Standard semi-CRF implementations materialize an O(TKC^2) edge tensor, limiting exact segment-level inference on long sequences. Flash-SemiCRF replaces that tensor with prefix-sum lookups and streaming forward-backward computation, while centered emissions stabilize scores and adapt duration preferences. The resulting fused Triton implementation makes exact structured inference viable at sequence lengths exceeding 10^6.

  • Problem

    Materializing O(TKC^2) edge tensors makes exact semi-CRF inference memory-bound and infeasible for long sequences.

  • Method

    Flash-SemiCRF uses prefix-sum edge evaluation, streaming dynamic programming, sublinear checkpointing, and centered emissions in a fused Triton kernel.

  • Results

    Exact semi-CRF inference reaches sequence lengths exceeding 10^6 by avoiding intermediate edge-tensor materialization.

  • Takeaways & Limitations

    Flash-SemiCRF makes coherent segmentations, explicit duration modeling, and segment-boundary uncertainty viable for long-range sequence modeling.

  • Takeaways & Limitations

    Boundary features are restricted to separable endpoint terms and cannot depend on duration or source label; gradient workspace remains quadratic in C.

Abstract

from arXiv · show

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.

1 Introduction

Semi-CRFs provide segment-level structure and uncertainty, but standard implementations become memory-prohibitive because they materialize dense edge potentials. Flash-SemiCRF avoids this bottleneck through on-the-fly scoring, streaming inference, checkpointing, and centered emissions.

  • Motivation: Semi-CRFs enforce valid partitions, model label-dependent durations, and provide posterior marginals over segment boundaries.These properties extend per-position neural predictions with segment-level structure and uncertainty.
  • Scaling barrier: O(TKC^2) edge tensors make standard semi-CRF inference memory-bound and infeasible for long genomic sequences.The cost depends on sequence length T, maximum duration K, and label count C.
  • Scaling barrier: Pruning methods are unsuitable for dense segmentation because every position belongs to a labeled segment and irrelevant-span sparsity does not generally hold.Thus, dense genomic tasks typically require traversing the full semi-CRF graph.
  • Flash-SemiCRF: Prefix-sum decomposition evaluates edge potentials on-the-fly in O(1) time per duration, eliminating the materialized edge tensor.The approach replaces stored edge potentials with compact cumulative scores.
  • Flash-SemiCRF: Ring-buffer streaming uses O(KC) working memory, while checkpointing provides sublinear memory growth in sequence length.These mechanisms preserve the forward-backward computation without retaining the full sequence of intermediate states.
  • Flash-SemiCRF: Globally centered emissions stabilize cumulative scores and induce an adaptive duration prior under label imbalance.The centering mechanism addresses numerical drift while changing duration preferences according to sequence-level emission statistics.

2 Terminological clarification

The paper distinguishes the model class from the algorithm used to compute it. A linear CRF fixes segment duration at one, whereas a linear scan denotes a sequential dynamic-programming implementation.

  • Terminology: A linear CRF is the K=1 model, where every segment has duration one.Its inference reduces to the standard per-position Viterbi or forward-backward recurrence.
  • Terminology: A linear scan is a left-to-right algorithm for executing the semi-CRF dynamic program.It can compute both linear CRFs and semi-CRFs with K ≥ 1.
  • Terminology: Linear CRF and linear scan are orthogonal distinctions between model structure and computational backend.A linear CRF can use different backends, and a linear scan can operate at multiple maximum durations.

3 Implementation

Flash-SemiCRF computes semi-CRF edge scores from centered prefix sums and streams dynamic programming instead of materializing the dense edge tensor. Its factorized boundary features improve streaming compatibility but restrict how boundary scores depend on segments.

  • Prefix-sum edge computation: O(TC) cumulative scores replace the O(TKC^2) edge tensor, enabling O(1)-time evaluation of each segment content score.The lookup uses differences of prefix sums for the segment interval.
  • Centered cumulative scores: The implementation subtracts a sequence-level emission baseline before constructing cumulative scores.This centered representation is used for subsequent segment-score lookups.
  • Centered cumulative scores: Centered accumulation reduces catastrophic cancellation when long sequences produce large cumulative scores with small segment-score differences.Without centering, cumulative scores grow as O(T), whereas centered sums have a smaller random-walk scale.
  • Edge construction: Each edge combines centered content, duration bias, transition, and boundary terms through constant-time array lookups.The source label, destination label, duration, and endpoint positions determine the required precomputed values.
  • Boundary features: Optional independent start and end projections add position-specific boundary scores to segment potentials.Separate projections allow segment entry and exit to use different encoder-derived signals.
  • Boundary features: The streaming factorization restricts boundary features to separable functions of destination label and one endpoint, excluding duration- or source-label-dependent boundary terms.Classical joint endpoint features require the full O(TKC^2) edge tensor.

4 Emission Baseline Centering

Per-label emission centering is both a numerical stabilization device and a modeling choice. It induces a sequence-adaptive duration prior that can counter label imbalance, especially for long segments of dominant labels.

  • Centering effect: Per-label centering does not cancel in the partition function because segmentations assign different labels to positions.This distinguishes it from path-invariant per-position centering.
  • Adaptive duration prior: The effective duration bias is Beff_k,c = Bk,c − νb,c · k, adapting duration preferences to sequence-level emission statistics.The centering contribution grows linearly with segment duration k.
  • Adaptive duration prior: Population-level duration biases combine with per-sequence centering adjustments, primarily regularizing sequence-specific variation under label imbalance.The interpretation is empirical-Bayes-like because Bk,c captures training-population effects while νb,c supplies sequence-level adjustment.
  • Worked example: Under label imbalance, centering penalizes long dominant-label segments while minimally penalizing or boosting rare labels.The worked example reports a −325 penalty for Intron and a +12 bonus for Promoter at k=100.
  • Ablation: Under 75/15/10% imbalance, mean centering recovered 59 rarest-label segments versus 19 with path-invariant centering, while dominant-label segments fell from 142 to 90.With balanced proportions, segment counts converged across centering modes.
  • Architectural context: Mean centering is suited to structured output decoders, whereas path-invariant centering is preferable when semi-CRF marginals become downstream features.The distinction preserves the canonical distribution in the downstream-feature setting.

5 Streaming Forward-Backward Algorithm

Flash-SemiCRF streams exact forward-backward inference by evaluating segment scores from prefix sums and retaining only bounded ring buffers plus sublinear checkpoints. Checkpoint normalization stabilizes long scans, while adaptive tiling and validation address numerical and GPU implementation constraints.

  • Streaming forward scan: A K-slot forward ring buffer stores only the K most recent messages because each recurrence looks back at most K positions.The buffer is indexed by t mod K, so the oldest entry is overwritten when it is no longer needed.
  • Normalization and checkpoints: Checkpoint-boundary normalization subtracts the current maximum, accumulates the shift, and saves the full buffer state with its cumulative normalizer.During backward recomputation, the saved normalizer restores the true scale needed for marginal computation.
  • Backward scan and the 2K ring buffer: The backward pass uses a 2K-slot ring buffer because writing β[t] otherwise overwrites the β[t+K] value still needed by the read window.This separates backward-pass reads and writes while preserving the bounded-memory scan.
  • Complexity: Working memory is O(KC), checkpoint storage is O(T/K · KC), and the prefix-sum array replaces the larger edge tensor.The resulting streaming dynamic-programming memory is sublinear in sequence length, while the prefix sums remain O(TC).
  • On-the-fly edge computation: Prefix-sum lookup replaces the O(TKC2) edge tensor, reducing dominant memory by a factor of KC while retaining O(TKC2) total time.At T = 10^6, K = 200, C = 6, the edge tensor requires ∼29 GB in float32, versus ∼24 MB for prefix sums.
  • Adaptive loop tiling: Adaptive destination-label tiling reduces register demand from approximately 384 to approximately 120 registers per thread, enabling 4–8 warps without spilling.The tile size is selected based on C to balance compile time, register pressure, and iteration count; iterations remain ≤8 at C = 256.
  • Correctness validation: Correctness is checked through probabilistic invariants, finite-difference gradients, and training convergence against independent implementations.At genome scale, invariant per-position deviation is 1.6×10^-4; training reaches relative final negative-log-likelihood difference <10^-9 with loss-curve cosine similarity 1.0.

6 Experiments

Experiments show that Flash-SemiCRF preserves stable scaling as state size and sequence length grow, avoiding the memory failures of tree-based alternatives. On TIMIT, it achieves similar phone error but better segmentation metrics and substantially faster training and inference.

  • Theoretical benchmarks: Tree-based methods degrade rapidly as (K+1)C increases, whereas linear and streaming scans maintain stable compute scaling.The fused Triton kernel matches this stability and improves throughput over unfused linear scans through on-the-fly edge evaluation.
  • Theoretical benchmarks: Linear and streaming scans maintain nearly constant memory as T grows, while tree-based and block-structured methods require orders of magnitude more memory.The fused kernel therefore maintains a nearly constant maximum state size at the out-of-memory frontier.
  • Genome-scale evaluation: At K = 2000–8000, fused Triton runtime reflects hardware-specific memory trade-offs, while competing methods are no longer tractable under the tested memory constraints.The reported trade-off is attributed likely to L1 cache saturation and may vary across systems.
  • TIMIT benchmark: On TIMIT, semi-CRF and linear CRF reach similar final PER, while semi-CRF improves boundary F1 and segment F1.Final PER is 0.218 versus 0.219; boundary F1 is 0.476 versus 0.468, and segment F1 is 0.215 versus 0.207.
  • Convergence: The semi-CRF starts behind but overtakes the linear CRF by epoch 5 on PER and by epoch 1 on boundary and segment F1.The PER gap peaks near epoch 20 at ∆≈0.005 and compresses to ∼0.001 at convergence.
  • Runtime: Flash-SemiCRF reduces TIMIT training time from 4,430 s to 175 s per epoch and inference time from 1,243 s to 7 s per pass.These correspond to 25× training and 178× inference speedups, with throughput of 6.5k versus 0.3k frames/s.
  • Boundary uncertainty: The semi-CRF reduces boundary entropy from 5.671 at initialization to 5.529 at epoch 50, concentrating probability mass near phoneme transitions.The effective boundary count drops from ∼295 to ∼252, while both models converge to Hpos = 3.651.
  • Per-utterance uncertainty: For one utterance, semi-CRF PER is 0.232 versus 0.411 for linear CRF, with boundary posterior peaks at true transitions and p < 0.2 inside segments.This illustrates posterior segmentation structure that the K=1 model cannot represent.

7 Discussion

Flash-SemiCRF addresses semi-CRF scaling by replacing edge-tensor materialization with prefix-sum computation and streaming dynamic programming. The approach improves scalability while exposing modeling trade-offs involving label count, maximum duration, numerical stability, and probabilistic fidelity.

  • Scaling and systems implications: Avoiding intermediate tensor materialization, rather than changing asymptotic complexity alone, reframes semi-CRF inference as compute-bound and supports sequences exceeding 10^6 positions.The method eliminates the O(TKC^2) edge tensor through prefix-sum decomposition and streaming dynamic programming.
  • Scaling and systems implications: Semi-CRFs become viable for long-range structured decoding, addressing per-position neural outputs that do not enforce globally consistent segmentations.The discussion connects scalable inference with segment-level modeling for long sequences.
  • Modeling consequences: Boundary and content gradients are structurally decoupled because start, end, and content terms enter the edge score additively.This block-diagonal Jacobian permits boundary and content scoring updates to adjust independently.
  • Modeling consequences: Centered emissions induce an adaptive duration prior that penalizes long segments of high-mean labels and relatively boosts low-prevalence labels.This can suppress degenerate frequent-state segmentations and improve recovery of rare classes in imbalanced settings.
  • Modeling consequences: The centered formulation trades numerical stability and regularization against model fidelity because it changes the underlying distribution and is not path-invariant.The choice between centered and uncentered scores therefore depends on whether strict probabilistic interpretation is required.
  • Limitations and implementation boundaries: The gradient workspace scales as O(B · Nckpt · K · C^2), so the implementation retains quadratic dependence on label count despite avoiding the edge tensor.The authors state that this is acceptable for genomic tasks typically using C ≤ 64, while maximum duration K remains a more consequential practical constraint.
  • Limitations and implementation boundaries: Segments longer than K lose unified duration and content scoring, making selection of K the primary tuning decision and motivating future analysis of degradation beyond K.The 95th-percentile intron length of approximately 22,908 bp approaches the regime where runtime and register pressure become significant.
  • Implementation boundaries: Streaming uses ring buffers with O(KC) working memory independent of sequence length, while checkpointing adds workspace proportional to B · Nckpt · K · C^2.The backward pass recomputes forward states from saved checkpoints, and deterministic workspace reductions avoid inter-segment atomic contention.

A.8 Implementation Correspondence

The appendix maps mathematical notation and backend provenance to implementations, contrasting tree-structured complexity with the Triton kernel’s on-the-fly edge computation.

  • The supplement defines n = (K + 1) · C as the expanded state-space size and Nckpt as the number of checkpoint segments.
  • All backends except the fused Triton kernel pre-materialize the O(TKC2) edge tensor, which dominates memory at large T.
  • The Triton kernel computes edge potentials on the fly, achieving true T-independent memory and enabling genome-scale inference.
  • The binary-tree backend computes semiring matrix products over the expanded state space with O(log T) parallel depth.
  • The tree backend requires O(T(KC)2) space, while checkpointing reduces memory with 2× compute overhead and extends viability from n < 100 to n < 150.

B.5 Backend 2: Block-Triangular Sparsity (Experimental)

Block-triangular and banded representations reflect local duration constraints but fail to preserve useful sparsity through tree composition, limiting their practical value for exact inference.

  • Block-Triangular Storage: The feasible duration pairs form a block-triangular pattern, so storing only valid dense C × C blocks can roughly halve storage.
  • Experimental Outcome: In PyTorch/CUDA experiments, block-triangular storage produced no speedups because indirect indexing and non-contiguous accesses outweighed skipping approximately 50% of blocks.
  • Structural Limitation: Tree backends use semiring matrix products over an expanded state space of size Θ(KC), where bounded duration does not prevent dense intermediate operators.
  • Structural Limitation: At a tree node, the duration compatibility matrix has anti-diagonal triangular structure rather than diagonal-banded structure.
  • Bandwidth Lower Bound: A clique among durations up to floor(S/2) forces bandwidth at least floor(S/2)C − 1 under every ordering.
  • Structural Limitation: When S ≥ 2K, the compatibility matrix is structurally dense because every pair of durations satisfies d1 + d2 ≤ S.
  • Fill-In Under Composition: Repeated composition widens reachability with bw(Bm) = min(T, mK), approaching dense width after logarithmically many tree levels.
  • Reordering: Reordering helps mainly at small spans: bandwidth ratios average approximately 0.44 for S ≤ K/2, 0.90 for K/2 < S ≤ K, and at least 0.97 for S ≥ K.

C.5 Practical Implications for Exact Semi-Markov Backends

Higher tree levels become near-dense through local compatibility and composition, so structural sparsity does not provide a reliable exact-inference optimization.

  • Higher tree levels dominate computation and memory because local duration compatibility and multi-step composition become near-dense.
  • Block-triangular formats reduce storage by roughly half, but indexing overhead negated this benefit at typical sparsity levels.
  • GPU sparse kernels generally underperform dense kernels unless sparsity exceeds approximately 90%.
  • Useful sparsity must therefore be introduced through approximation methods such as segment filtering or hybrid architectures.

Appendix D Linear and Near-Linear CRF Implementations

Appendix D provides specialized K = 1 and K = 2 paths, boundary mechanisms, centered cumulative scores, and marginal-gradient details for streaming CRF implementations.

  • Dispatch: The implementation dispatches automatically by K; full boundary projections force the generic K ≥3 path, while scalar boundaries do not affect dispatch.
  • Score Centering: Centering subtracts a masked per-sequence baseline before accumulation, bounding cumulative-score magnitude without changing per-position emission differences.
  • Boundary Handling: Scalar sequence boundaries are folded into cumulative scores with prefix-sum arithmetic, leaving the forward and backward algorithms unchanged.
  • Boundary Handling: Position-dependent boundary projections provide richer start and end preferences but require the generic streaming path.
  • Specialized Paths: K = 1 reduces the Semi-CRF to a standard linear-chain CRF, while K = 2 supports durations one and two through explicit history variables.
  • Complexity: The K = 1 path uses O(TC2) time and stores full α history for backward computation, while K = 2 uses two matrix-vector products per timestep.
  • Marginals: Reference marginal computation clamps log marginals to [−80, 80] before exponentiation to prevent overflow and underflow.

D.10 Implementation Summary

The implementation automatically selects the appropriate streaming implementation based on maximum segment duration. Implementation characteristics are organized by maximum segment duration.

  • Dispatch automatically selects the appropriate streaming implementation based on K.Users do not need to manage dispatch manually.
  • The supplied implementation summary identifies K as the dispatch criterion and maximum segment duration as the table’s organizing dimension.
  • Table D.3 organizes implementation characteristics by maximum segment duration.

D.11 Functional Equivalence

The K = 1, K = 2, and K ≥3 implementations compute the same partition function and gradients under their respective duration constraints. Specialized paths optimize performance rather than behavior, with equivalence conditions for full boundary projections explicitly constrained.

  • All three implementations compute the same partition function and gradients for their respective segment duration constraints.
  • K = 1 and K = 2 fast paths match their corresponding general streaming kernels when full boundary projections are inactive.Scalar boundaries are folded into S and do not affect equivalence.
  • Specialized paths exist purely for performance optimization, not behavioral differences.
Loading 2604.18780v1…