Technical Summary: Approximate Structured Diffusion for Sequence Labelling

Problem Statement

Sequence labelling is a fundamental Natural Language Processing (NLP) task involving the assignment of a label to each token in an input sentence (e.g., Part-of-Speech tagging, Named Entity Recognition). While modern approaches often utilize Linear-Chain Conditional Random Fields (CRFs) parametrized by neural networks, these models face inherent limitations. Standard CRFs typically restrict interactions to adjacent labels (bigrams) to ensure tractability, which limits their expressivity when modeling long-range dependencies. Conversely, while diffusion models have shown promise in language modeling by conditioning generation on unbounded contexts, applying them to structured prediction tasks like sequence labelling presents challenges. Specifically, exact inference in CRFs (e.g., via Viterbi or Forward-Backward algorithms) is computationally expensive and difficult to parallelize, making the iterative sampling required by diffusion models prohibitively slow.

Methodology

The authors propose a novel framework that bridges structured prediction and discrete diffusion to train a CRF conditioned on an entire label sequence, specifically a noisy version of that sequence.

1. Discrete Diffusion Framework

The model treats sequence labelling as a generation task where a clean label sequence $y_0$ is corrupted by a forward diffusion process $q$ over $T$ steps to reach a random sequence $y_T$. The goal is to learn a denoiser $p_\theta$ that reverses this process, predicting the clean sequence from noisy inputs.

• Forward Process: At each step $t$, noise is added to the labels according to a corruption ratio $\beta_t$, eventually resulting in a uniform distribution over all possible label sequences.
• Denoising Process: The model learns to predict the clean sequence $y_0$ (or its expected value $\hat{y}_0$) given a noisy sequence $y_t$ and the input sentence $s$.

2. Structured Denoising with CRFs

Unlike standard diffusion models that predict tokens independently, this approach implements the denoiser $p_\theta$ using a CRF.

• Conditioning: The CRF is conditioned not only on the input sentence but also on the noisy label sequence. This allows the model to consider unbounded label interactions while maintaining the ability to enforce preferences on adjacent labels.
• Approximate Inference: Since exact sampling from a CRF distribution is costly ($O(n)$ complexity per step) and hinders parallelization, the authors approximate the CRF distribution using Mean-Field inference. This allows for position-wise independent computation, recovering the efficiency of unigram models while retaining structural benefits.
• Training Objective: The model is trained by maximizing a variational lower bound of the log-likelihood. The loss function involves minimizing the Kullback-Leibler (KL) divergence between the true posterior (conditioned on the gold sequence) and the model's posterior (conditioned on the denoised sequence). A denoising loss (negative log-likelihood) is also added to improve stability.

3. Neural Architecture

The architecture employs a Diffusion Transformer (Peebles and Xie, 2022) for the denoising step:

• Input: Noisy label embeddings (from a trainable lookup table) and contextual word embeddings (from a Transformer encoder or CharLSTM).
• Context: The timestep $t$ is embedded and concatenated with word embeddings to guide the normalization.
• Output: The Transformer blocks output vectors that are converted via MLPs into unigram scores and, for structured models, pairwise transition scores (bigrams).
• Decoding: The decoding process uses a halving strategy, jumping from step $t$ to $\lfloor t/2 \rfloor$ iteratively until reaching step 0, reducing the number of denoiser calls to be logarithmic in $T$.

Key Contributions

1. Integration of Diffusion and CRFs: The paper introduces a method to train a CRF conditioned on a noisy label sequence, effectively combining the global context capabilities of diffusion with the structural constraints of CRFs.
2. Approximate Structured Inference: To address the computational bottleneck of CRF sampling in diffusion loops, the authors utilize Mean-Field approximation. This enables efficient, parallelizable training and decoding while preserving the ability to model long-range dependencies.
3. Scalability: The proposed method demonstrates superior scaling properties compared to baseline CRFs. While traditional CRFs often plateau or overfit as parameter counts increase, the diffusion-based models continue to improve with larger architectures.

Experimental Results

The authors evaluated their approach on four datasets from Universal Dependencies v2.15 (English, German, French, and Dutch) for Part-of-Speech (POS) tagging.

• Performance: The structured diffusion model with Mean-Field approximation (Diffusion-MF) achieved the best results across all datasets. Compared to the best non-diffusion baseline (standard CRF), the proposed method achieved an average 16.54% error reduction.
• Comparison:
  • Diffusion-MF outperformed Diffusion-Uni (unstructured diffusion), indicating that the added CRF structure improves performance over simple diffusion.
  • Both diffusion variants outperformed the three baselines: Unigram, Mean-Field, and standard CRF.
• Scaling: Experiments varying model sizes (from XS to XL) showed that diffusion models scale effectively with increased parameters, whereas baseline models showed diminishing returns or performance degradation after a certain point.
• Efficiency: While the structured approach (Diffusion-MF) is slower during evaluation than unstructured diffusion due to the iterative denoising process, it remains comparable in training speed to standard CRFs. The authors note that exact structured diffusion (Diff-CRF without approximation) is prohibitively slow, validating the necessity of the Mean-Field approximation.

Significance and Claims

The paper claims that this approach offers a new paradigm for sequence labelling that reconciles the need for long-range evidence with local constraints. By placing a CRF inside a diffusion denoiser loop, the model achieves:

1. Guaranteed Normalization: Ensuring valid probability distributions.
2. Global Context: Leveraging the evolving noisy sequence to inform predictions.
3. Parallelizability: Achieving efficiency through Mean-Field updates.

The authors conclude that while the method incurs higher computational costs (memory and energy) compared to simpler models, it provides a robust solution for tasks requiring complex label interactions. They suggest the approach is applicable to other NLP tagging tasks such as Named Entity Recognition (NER) and word segmentation. The paper modestly acknowledges that the increased performance comes at the expense of energy consumption and compute time, and notes that the convergence of the parallel Mean-Field approximation is not theoretically guaranteed, though it performed well empirically.