1 Introduction↩︎

Sequence labelling, a fundamental task in NLP, consists in assigning a tag to each token of an input sentence. It is the foundation of a variety of NLP applications, such as POS tagging, named-entity recognition or parsing.

Modern approaches to this task are based on CRFs parametrised by neural networks [1], [2]. While structured models such as CRFs consider interactions between labels, tractability impose some restrictions. Thus many proposed models are limited to bigrams, i.e.. correlations of adjacent labels.

Recently, diffusion has been applied to language modelling to effectively condition generation on unbounded contexts [3]–[5]. In practice, these models train denoisers to predict independently each token from noisy versions of the clean output.

In this work, we bridge these two concepts, i.e.. structured prediction and discrete diffusion, for sequence labelling. We define a CRF conditioning the predicted label sequence not only on the input sentence but also on a noisy label sequence. This helps the model consider unbounded label interactions while remaining able to enforce preferences on predicted adjacent labels. Decoding with diffusion models requires iterative sampling to refine predictions from random noise. Since sampling CRF distributions is costly, with a complexity linear in the input size, we speed up decoding and training by approximating them with Mean-Field.

We evaluate on POS tagging and show that this model scales up better than baseline CRFs, with both the unigram diffusion model and the addition of either a CRF denoiser or its Mean-Field approximation achieving superior performance.

2 Model↩︎

2.1 Standard Sequence Labelling Model↩︎

Given a sentence \(\boldsymbol{s}= s_{1}\dots s_{n}\), with \(s_{i}\) the \(i^{\text{th}}\) word, labelling produces a sequence \(\boldsymbol{y}=y_{1}\dots y_{n}\), with \(y_{i}\in \mathcal{L}\) the label for \(s_{i}\). More precisely, sequence labelling models define a parametrised probability distribution \(p_\theta (\boldsymbol{y} | \boldsymbol{s})\) so labelling amounts to returning the mode \(\widehat{\boldsymbol{y}} = \mathop{\mathrm{argmax}}_{\boldsymbol{y}} p_\theta (\boldsymbol{y} | \boldsymbol{s})\) and learning parameters \(\theta\) is cast as Maximum Likelihood Estimation. These distributions are usually written as energy models \(p_\theta (\boldsymbol{y} | \boldsymbol{s}) \propto \exp f(\boldsymbol{x}, \boldsymbol{y};\theta)\), computed by a neural network implementing \(f\), i.e.. parameters \(\theta\) are the parameters of \(f\). The decomposition of \(f\) over sequences is crucial for efficiency.

Unigram Models sum unary potential over the sequence \(f(\boldsymbol{s}, \boldsymbol{y};\theta) = \sum_{i=1}^{n} f(\boldsymbol{s}, y_i;\theta)\). Typically this is implemented as a Transformer [6] whose output vector at position \(i\) feeds a MLP computing \(f\) for all labels at this position. As a consequence of the decomposition of \(f\), \(p_{\theta}\) is factorized, i.e.. \(p_{\theta}(\boldsymbol{y} | \boldsymbol{s}) = \prod_{i=1}^{n} p_{\theta}(y_i | \boldsymbol{s})\). Labelling and training are efficient but the independence between predictions impairs the expressivity required to model fine-grained label interactions.

Bigram Models sum unary and binary² potentials over adjacent positions, \(f(\boldsymbol{s}, \boldsymbol{y};\theta) = \sum_{i=1}^{n} f_1(\boldsymbol{s}, y_i;\theta) + \sum_{i=1}^{n-1} f_2(\boldsymbol{s}, y_i, y_{i+1};\theta)\). Transformer output are fed to \(f_1\) as in the unigram case. For \(f_2\), we need to compute a transition matrix for all labels \(y_i, y_{i+1}\). This can be either implemented as a position-independent matrix, or as a MLP computing at each position the transition weights fed by a Transformer’s output. Labelling and training can be performed in linear time w.r.t.. sentence’s length, as can be computed marginal probabilities, with Viterbi [7] and Forward/Backward [8] algorithms. Bigram CRFs are difficult to parallelise but approximations such as Mean-Field [9] or Mean-Regularisation [10] recover the position-wise independent computation, and thus the efficiency, of unigram models.

2.2 Discrete Diffusion Models for Labelling↩︎

We follow diffusion language models [3], [4] to define labelling as generation of labels given words.

Forward Diffusion. A sequence of tags \(\boldsymbol{y}^{0}\) is altered by a forward diffusion process \(q\) consisting of \(T\) steps \(\mathbf{y}^{1}\dots \mathbf{y}^{T}\) to eventually obtain a random sequence³ \(\boldsymbol{y}^{T}\). Generating such sequences is a Markovian process \(q(\mathbf{y}^{1}\dots \mathbf{y}^{T} | \mathbf{y}^{0}) = \prod_{t=1}^T q_{t} (\mathbf{y}^{t} | \mathbf{y}^{t-1})\) with independent noise at each position \(i\): \(q_{t} (\boldsymbol{y}^{t} | \boldsymbol{y}^{t-1}) = \prod_{i=1}^n q_t(y_i^t | y_{i}^{t-1})\).

Noise distributions are parametrised by a corruption ratio \(\beta_{t}\) following a predefined schedule:⁴ \[\label{eq:noise-step} q_{t}(y_{i}^{t+1}| y_{i}^{t}) = \begin{cases} \beta_{t} + \frac{1-\beta_{t}}{|\mathcal{L}|} \text{ if } y_{i}^{t+1} = y_{i}^{t}\\ \frac{1-\beta_{t}}{|\mathcal{L}|} \text{ otherwise} \end{cases}\tag{1}\] The parameters of these conditional distributions \(q_t (\cdot|\cdot)\) for timestep \(t\) can be encoded as a matrix \(Q_{t}\). We can also precompute consecutive applications of \(t\) diffusion steps \(q_{0|t} (y_{i}^{t} | y_{i}^{0}) = \sum_{y_{i}^{t-1}} q_{t}(y_{i}^{t}|y_{i}^{t-1}) q_{0|t-1}(y_{i}^{t-1} | y_{i}^{0})\).

Denoising. Our model, following [3], [4] for language models, produces a parameterised distribution on label sequences from a random sequence by reversing the diffusion process. With a slight abuse of notation we also denote this distribution as \(p_{\theta}\). This model can assign a probability to less and less noisy sequences, also a markovian process : \(p_{\theta} (\mathbf{y}^0\mathbf{y}^{1}\dots \mathbf{y}^{T}| \boldsymbol{s})= p(\mathbf{y}^{T}) \prod_{i=t}^T p_{\theta}(\mathbf{y}^{t-1}|\mathbf{y}^{t}, \mathbf{s})\), where the prior \(p(\boldsymbol{y}^{T})\) is the uniform distribution. We drop the condition on \(\boldsymbol{s}\) in notations.

Denoiser \(p_{\theta}\) is implemented by a neural network presented in §2.3. The same network is used for all \(t\): to add time information, we feed the neural network with a learned representation of \(t\). We follow the widely adopted architecture of [11] and describe a single denoising step from \(t\) as full denoising followed by \((t-1)\) forward steps:

\[\begin{align} p_{\theta}(\boldsymbol{y}^{t-1}|\boldsymbol{y}^{t}) &= \sum_{y^{0}} p_{\theta}(\boldsymbol{y}^{0}|\boldsymbol{y}^{t}) q(\boldsymbol{y}^{t-1} | \boldsymbol{y}^{t}, \boldsymbol{y}^{0})\\ &= \mathbb{E}_{\boldsymbol{y}^{0}\sim p_{\theta}(\cdot|\boldsymbol{y}^{t})} \big[ q(\boldsymbol{y}^{t-1} | \boldsymbol{y}^{t}, \boldsymbol{y}^{0}) \big]\\ &\approx q(\boldsymbol{y}^{t-1} | \boldsymbol{y}^{t}, \widehat{\boldsymbol{y}^{0}} )\\ &\text{ with } \widehat{\boldsymbol{y}^{0}}= \mathbb{E}_{\boldsymbol{y}^{0}\sim p_{\theta}(\cdot|\boldsymbol{y}^{t})}[\boldsymbol{y}^{0}]. \end{align}\]

A denoising step can thus be modelled as sampling from \((i)\) the so-called posterior distribution with \((ii)\) the clean sequence \(\boldsymbol{y}_0\) replaced by an expected sequence \(\widehat{\boldsymbol{y}^{0}}\). In practice, addressing \((i)\) requires computing the posterior distribution, expressed with three tractable distributions, from Bayes’ theorem and Markovian assumption :

\[\begin{align} q(\boldsymbol{y}^{t-1} | \boldsymbol{y}^{t}, \boldsymbol{y}^{0}) &= \frac{q( \boldsymbol{y}^{t-1}, \boldsymbol{y}^{t} | \boldsymbol{y}^{0}) }{q(\boldsymbol{y}^{t}| \boldsymbol{y}^{0})} \\ &= \frac{q( \boldsymbol{y}^{t} | \boldsymbol{y}^{t-1} , \boldsymbol{y}^{0}) q( \boldsymbol{y}^{t-1} | \boldsymbol{y}^{0}) }{q(\boldsymbol{y}^{t}| \boldsymbol{y}^{0})} \\ &= \frac{q_{t}( \boldsymbol{y}^{t} | \boldsymbol{y}^{t-1}) q( \boldsymbol{y}^{t-1} | \boldsymbol{y}^{0}) }{q(\boldsymbol{y}^{t}| \boldsymbol{y}^{0})}. \end{align}\]

While explanations above indicated that denoising is performed step by step from \(t\) to \(t-1\), we can rewrite it to perform several steps at once, from \(t\) to \(t-k\). This may impact the quality of the generated sequence since the denoiser is called less, and has thus less opportunities to rely on the input sequence. In our experiments we use a halving strategy and go from step \(t\) to step \(\lfloor \frac{t}{2} \rfloor\), starting with step \(T\) until we reach \(0\), so the number of calls to the denoiser is logarithmic in the number of diffusion steps.

Structured Denoising We can adapt the previous decoding method to the case where the denoiser \(p_{\theta}\) is implemented by a CRF. Remember that the denoiser’s role is to generate \(\widehat{\boldsymbol{y}^{0}}= \mathbb{E}_{\boldsymbol{y}^{0}\sim p_{\theta}(\cdot|\boldsymbol{y}^{t})}[\boldsymbol{y}^{0}]\) the fractional counts of each variable, given by the marginals probabilities of labels. For linear-chain CRFs we can compute marginals in \(O(n)\) time complexity, either by running the forward-backward algorithm [8] or backpropagating through the log-partition [12]. Unfortunately, this approach is intractable in our context because of the limited parallelisability of the Viterbi algorithm or its variants. Moreover the denoiser must be called multiple times at decoding and the linear space complexity of these methods also burdens training with memory consumption.

Instead, we can approximate the CRF distribution with Mean Regularisation [10] or find the closest factorised distribution with Mean-Field [9]. We experiment with the latter and show we can exploit structures with diffusion models while remaining efficient.

Training is performed by maximizing likelihood with the denoiser synced to the diffusion model at each timestep. More precisely we optimise a variational lower bound of the log-likelihood :

\[\begin{align} &\log p_{\theta}(\boldsymbol{y}^{0}) \nonumber = \log \sum_{y^{1}\dots y^{T} \sim q(\cdot | y^{0})} p_{\theta}(\boldsymbol{y}^{0},\boldsymbol{y}^{1},\dots,\boldsymbol{y}^{T}) \nonumber \\ &\ge \mathop{\mathrm{\mathbb{E}}}_{\boldsymbol{y}^{1}\sim q_{0|1}(\cdot\mid \boldsymbol{y}^{0})} \Big[\log p_{\theta}(\boldsymbol{y}^{0}\mid \boldsymbol{y}^{1})\Big] \nonumber\\ &\quad-\sum_{t=2}^{T} \;\mathop{\mathrm{\mathbb{E}}}_{\boldsymbol{y}^{t}\sim q_{0|t}(\cdot\mid \boldsymbol{y}^{0})} \left[\, D_{\mathrm{KL}}\!\vcenter{ \left( \begin{array}{@{}l@{}} q(\boldsymbol{y}^{t-1}\mid \boldsymbol{y}^{t},\boldsymbol{y}^{0}) \\ [.6ex] || \; p_{\theta}(\boldsymbol{y}^{t-1}\mid \boldsymbol{y}^{t}) \end{array} \right)}\,\right] \nonumber\\ &\quad- D_{\mathrm{KL}}\Big(q_{0|T}(\boldsymbol{y}^{T}\mid \boldsymbol{y}^{0}) \,\big\|\, p(\boldsymbol{y}^{T})\Big), \end{align}\] where \(D_{KL}\) is the Kullback-Liebler divergence between the two distributions. The last term can be ignored since, by definition, the two distributions are uniform and their divergence is thus zero. We train our model by uniformly sampling a \(t\) between 1 and \(T\), and then a sequence \(\boldsymbol{y}^{t}\sim q_{0|t}(\cdot|\boldsymbol{y}^{0})\), which simulates \(t\) diffusion steps. If \(t=1\), we only consider the first term; si \(t\geq 2\) we only consider the mean in the second term based on this \(t\), which brings us back to optimizing a single KL divergence.

We notice that once sampled \(\boldsymbol{y}^{t}\), the first term is just a log-likelihood. For the second term, we get:

\[\begin{align} &D_{KL}[q(\boldsymbol{y}^{t-1} | \boldsymbol{y}^{t}, \boldsymbol{y}^{0})||p_{\theta}(\boldsymbol{y}^{t-1}|\boldsymbol{y}^{t})] \\ =&D_{KL}[q(\boldsymbol{y}^{t-1} | \boldsymbol{y}^{t}, \boldsymbol{y}^{0})||q(\boldsymbol{y}^{t-1} | \boldsymbol{y}^{t}, \widehat{\boldsymbol{y}^{0}})] + C, \end{align}\] i.e.. we seek to match posteriors conditioned on respectively gold sequences and the denoised ones.

Finally, we add the denoising loss [4], i.e.. the negative log-likelihood given by the denoiser to training instances, as it helps training stability and convergence.

2.3 Neural Architecture↩︎

The neural architecture depicted in Fig.1 implements potential functions, \(f\) for unigrams or \(f_1,f_2\) for bigrams, as defined in §2.1, which are then summed to define probability distributions. All models start converting words \(s_1,\dots,s_n\) to non-contextual representations with a look-up table and a charLSTM [13]. These are contextualized with Transformers [6] to obtain vectors \(e_1,\dots,e_n\). Alternatively, pretrained embeddings can be used. The unigram model is parametrised by \(n\) vectors \(l_i\) of \(|\mathcal{L}|\) scores computed by a MLP from \(e_i\). The bigram models adds \(|\mathcal{L}|\times|\mathcal{L}|\) scores which represent label transitions from one position to the next. These are computed for each position by a MLP.

Our diffusion models implement the denoiser with Diffusion Transformer blocks [14]. It takes as input a sequence of label embeddings from a trainable look-up table corresponding to noisy labels and, as context for normalization, the concatenation of the contextual word embeddings and a trainable embedding of the timestep. After these blocks a final MLP converts position vectors to unigram, and possibly bigram, scores.

3 Experiments↩︎

Data We experiment on 4 datasets from Universal Dependencies v2.15 [15] namely EN-EWT, DE-GSD, FR-GSD, NL-LassySmall for English, German, French, and Dutch. We use standard splits, evaluate accuracy with punctuation and average results over 8 random seeds.

We compare the 3 baseline models, unigram, Mean-Field and CRF with our diffusion models, the unstructured and structured models, Diffusion-Uni and Diffusion-MF respectively. Additionally, we report the best POS taggging results from [10]. The hyper-parameter setups for each model are described in 12, while additional results, with a different encoder, features and a baseline for Diff-CRF (the non approximated structured diffusion model) in 8, an ablation study in 9, and the models’ speeds in 11.

Our results in Table ¿tbl:tab:ud95dev? show that structured diffusion improves performances in nearly all datasets, and get an average \(16.54\%\) error reduction between the best model without diffusion (CRF) and the model with structured diffusion (Diffusion-MF). We find that the diffusion approach allows models to scale better with more parameters, surpassing the 3 baseline with an equal parameter count, this is shown in Table 9 in Appendix. Moreover in this setting, while the main source of performance increase is the use of diffusion, we can see that the additional structure given by the CRF (Diffusion-MF) improves performance over the simple diffusion process (Diffusion-Uni).

4 Related Work↩︎

Structured discrete diffusion. Discrete diffusion either corrupts labels directly [3], [4] or relaxes tokens for Gaussian diffusion and discretise at decode time [11], [14], [16]; related variants use iterative masking [17], bit encodings [18], or VQ codebooks [19]. For structured NLP, diffusion has been applied to span-level NER [20], token-level labelling via bit-relaxed sequences [21], and non-autoregressive constrained generation [22]; continuous-time categorical formulations connect to jump processes and score matching [23]. Adjacent works insert CRFs around diffusion, e.g.. continuous CRF for latent diffusion [24] or diffusion-enhanced BiLSTM-CRF for NER [25]. Unlike previous works, we use a CRF denoiser inside the loop, yielding (i) guaranteed normalization, (ii) global context via the evolving noisy sequence, and (iii) parallelizable mean-field updates.

Structured prediction with neural CRFs. Neural CRF models remain strong for labelling [1]; differentiable inference via unrolled mean-field (CRF-as-RNN) and parallel accelerations improve efficiency [2], [9], [10], alongside amortized perspectives [26]–[28] and classical variational analyses [29]–[31]. Rather than single-shot CRF decoding, we perform repeated, globally informed CRF diffusion-driven updates, reconciling long-range evidence with local constraints beyond purely accelerated CRF decoders [9], [10]. Our model is then closer to [32] but with the CRF inside a diffusion denoiser instead of generative Transformers.

5 Conclusion↩︎

We presented a novel approach to sequence labelling, with application to POS tagging, based on structured prediction and discrete diffusion to better predict tag distributions. Our model improves the tagging metric and increases parameter scaling, surpassing the baseline and performing even better as the number of parameters grows while keeping a manageable time complexity thanks to the Mean-Field approximation of the CRF distribution. Our approach could also be applied in other tagging tasks in NLP, such as NER, or word segmentation.

Limitations↩︎

We showed that the presented method can scale, i.e.. the more parameters the better accuracy is, as opposed to prior methods which tend to overfit when the number of parameters grow. However, this increase come at the expense of memory consumption and compute time. In other words, our models require more energy to be run at their full potential.

In order to improve parallelization, we resort to the parallel version of Mean-Field for which convergence is not guaranteed. Although we didn’t see pathological divergence in practice, we note that the method recently developed by [10] could be used as a drop-in replacement for parallel Mean-Field with convergence guarantees. TODO: Change if we need B-CRF for NER or POS-tagging.

Ethical Considerations↩︎

We believe that our work does not raise ethical concerns. We present a novel architecture for sequence labelling based on diffusion and structured prediction and we test it on standard, publicly available data.

We acknowledge the environmental impact of the energy cost of training neural models.

6 NER results↩︎

7 Decoding Algorithms↩︎

Algorithm 2 describes the unbatched version of the decoding process with timesteps halved at each round. It relies on five subroutines:

1. sample samples independently at each position of the input sentence according to its input categorical distribution;

2. tag-embed and time-embed embed discrete values (resp. POS labels and timesteps) to their dense representations. They are implemented as look-up tables;

3. denoise-marginals calls the denoiser on the input sentence \(x\) from a noisy prediction \(\widetilde{y}\) and compute marginal distributions of all labels at each positions. For unstructured diffusion or Mean-Field, we simply return the softmaxed logits return by the neural network implementing the denoiser. For structured diffusion we use the forward/backward algorithm.

4. posterior computes the posterior distribution from marginals \(p\) between timesteps \(t\) and \(t'\).

The algorithms starts by sampling labels from the uniform distribution. Then at each relevant timestep, it performs the following operations. First the previously predicted labels, as well as the current timestep \(t\), are embedded. Then the denoiser is called to predict a new sequence of marginal label distributions. Finally, we use the posterior distribution to sample labels at timestep \(t'=\left\lfloor \frac{t}{2} \right\rfloor\). These steps are repeated recursively for timestep \(t'\), until we reach timestep 1.

8 Additional Results↩︎

In Table ¿tbl:tab:glove95testing?, we report the results obtained using the GloVe pretrained embeddings [33], which strengthen the validity of our approach, even when using more advanced pretrained features.

In Tables ¿tbl:tab:ud95dev95diff80M? and ¿tbl:tab:ud95dev95diff650M?, we compare the best baseline setups for the unigram, mean-field and CRF taggers, to the L and XL versions of our models (see 12) respectively using either a transformer as an encoder, or a pre-trained multilingual transformer model, RoBerta-Large [34]. Note that due to training time constraints, we only trained Diffusion-CRF, our structured diffusion model which doesn’t use the mean-field approximation using RoBerta-Large as an encoder. Model sizes are described in 12, while a global view of the scaling, or lack thereof for each model is presented in 9. The 2 tables show that both structured and unstructured diffusion work well, however, we acknowledge that the bitter lesson applies here, where with an the increased parameter count of the XL model, as well as with the enriched word embeddings with RoBerta-Large, the unstructured model manages to catch up to Diff-MF.

9 Ablation Studies↩︎

We find that the diffusion approach allows models to scale better with more parameters, surpassing the 3 baseline models when using a low parameter setup, but also getting increasingly better results as the parameter count increases, while the baseline models show no such improvements passed a very low ceiling after which performance actually decreases. Given that the baseline models consist principally of an encoder and a scoring MLP, we only tested their scaling when using a transformer as an encoder, whereas with the diffusion models, the encoder’s size remains the same, but the parameters allotted to the diffusion and denoising processes vary, this is further explained in 12. The principal limit of our scaling, to our knowledge is the limited memory in the GPUs we used.

Another test we did was to use different mean-field iterations, which led to increasing though diminishing gains until 10 iterations for the baseline mean-field model, and 15 iterations for Diff-MF. The effect of different numbers of iterations on speed is shown in 11.

10 Diffusion Tuning↩︎

We tested different numbers of layers for the diffusion transformer and settled on 8, which yields high results in all tested datasets while having a memory footprint small enough to ensure that both the training and evaluation can be carried out with no errors. As can be seen in ¿tbl:tab:layers95testing?, the unstructured model keeps getting better results with added layers, and more parameters.

11 Timing Experiments↩︎

In this section we notice in that a high number of mean-field iterations as well as the diffusion itself both have a non negligible effect on speed in Table ¿tbl:tab:mf95iter95timing?, however it is worth noting that the slowest training speeds we obtained for Diff-MF were about equivalent with the speed of the CRF model. Where our structured approach does slow down is in the evaluation, due to our decoding needing to perform the denoising process multiple times. Diff-CRF in particular is prohibitively slow, further validating our approximation approach with mean-fields. Similarly, our diffusion models do slow down considerably the more layers we use, as shown in ¿tbl:tab:layers95timing?.

12 Hyperparameters↩︎

We categorize our models into 5 different size configurations, based on the parameter count they have for the main baseline, XS, S, M, L, XL, corresponding to roughly 20, 40, 60, 80 and 650 million parameters respectively. 20 million parameters are allotted to the transformer encoder, for the baseline models, unigram, mean-field, CRF, the encoder is given more parameters to reach the bigger configurations, as the encoder is all these models have, save for a scoring MLP. For the diffusion models, the encoder does not change in size, but we change the size of the diffusion transformer used, thus in the M configuration for example, the non diffusion models have 60 million parameters allocated to their encoder, while the diffusion models have 20 million for the encoder, and 40 million for the diffusion transformer. We thus cannot have a XS diffusion baseline, as it would have 0 parameters for the diffusion, and we also decided not to test out the XL baselines for the non diffusion models, as none of them showed potential to perform any better than in the smaller baselines. All models use no dropout for their scoring MLP. Below is a few list of vector sizes, layer counts and the formula which gives the parameter count of the models.

References↩︎