Introduction

A Conditional Random Field is useful for sequence labeling tasks where neighboring tags influence each other, such as named entity recognition or part-of-speech tagging. In TensorFlow 2, a practical approach is to let a neural network produce token-level scores and then train and decode those scores with the CRF utilities from TensorFlow Addons.

Why Add a CRF Layer

A plain softmax classifier predicts each token independently. That is often good enough, but it can produce inconsistent label sequences. For example, a tagging scheme may allow I-PER only after B-PER, yet an independent classifier can still output invalid transitions.

A CRF learns transition scores between labels and decodes the best global path. That usually improves sequence consistency even when the encoder is unchanged.

Model Structure

A common TensorFlow 2 setup looks like this:

• embedding layer
• sequence encoder such as BiLSTM
• dense layer producing emission logits
• CRF log-likelihood for training
• Viterbi decode for inference

The neural network predicts emission scores for each token and tag. The CRF adds transition scores between successive tags.

Minimal TensorFlow 2 Example

The model below produces emission logits. The CRF transition matrix is stored separately so it can be passed to the Addons functions.

1import tensorflow as tf
2import tensorflow_addons as tfa
3
4
5class BiLstmCrf(tf.keras.Model):
6    def __init__(self, vocab_size, emb_dim, hidden_dim, num_tags):
7        super().__init__()
8        self.embedding = tf.keras.layers.Embedding(
9            input_dim=vocab_size,
10            output_dim=emb_dim,
11            mask_zero=True,
12        )
13        self.encoder = tf.keras.layers.Bidirectional(
14            tf.keras.layers.LSTM(hidden_dim, return_sequences=True)
15        )
16        self.classifier = tf.keras.layers.Dense(num_tags)
17        self.transitions = tf.Variable(
18            tf.random.uniform(shape=(num_tags, num_tags)),
19            name="crf_transitions",
20        )
21
22    def call(self, token_ids, training=False):
23        x = self.embedding(token_ids)
24        x = self.encoder(x, training=training)
25        return self.classifier(x)

The output shape is batch size by sequence length by number of tags. Those are the emission scores consumed by the CRF functions.

Compute Sequence Lengths and Loss

Padding is common in sequence batches, so CRF training also needs the true sequence lengths:

1def sequence_lengths(token_ids):
2    mask = tf.cast(token_ids != 0, tf.int32)
3    return tf.reduce_sum(mask, axis=1)
4
5
6@tf.function
7def train_step(model, optimizer, token_ids, tag_ids):
8    lengths = sequence_lengths(token_ids)
9
10    with tf.GradientTape() as tape:
11        logits = model(token_ids, training=True)
12        log_likelihood, _ = tfa.text.crf_log_likelihood(
13            logits,
14            tag_ids,
15            lengths,
16            transition_params=model.transitions,
17        )
18        loss = -tf.reduce_mean(log_likelihood)
19
20    variables = model.trainable_variables + [model.transitions]
21    gradients = tape.gradient(loss, variables)
22    optimizer.apply_gradients(zip(gradients, variables))
23    return loss

This training step uses tfa.text.crf_log_likelihood, which is the core operation for linear-chain CRF training in TensorFlow Addons.

Decoding with Viterbi

During inference, you want the highest-scoring tag sequence, not just the per-token argmax. Use crf_decode:

1def decode_tags(model, token_ids):
2    lengths = sequence_lengths(token_ids)
3    logits = model(token_ids, training=False)
4    decoded, scores = tfa.text.crf_decode(
5        logits,
6        model.transitions,
7        lengths,
8    )
9    return decoded, scores

The returned tag IDs represent the globally best path under the learned transition rules.

End-to-End Dummy Run

This small example shows the pieces working together on synthetic data:

1vocab_size = 50
2num_tags = 4
3
4model = BiLstmCrf(vocab_size=vocab_size, emb_dim=16, hidden_dim=8, num_tags=num_tags)
5optimizer = tf.keras.optimizers.Adam(1e-3)
6
7token_batch = tf.constant([
8    [3, 8, 2, 0, 0],
9    [4, 6, 7, 9, 1],
10], dtype=tf.int32)
11
12tag_batch = tf.constant([
13    [1, 2, 3, 0, 0],
14    [1, 1, 2, 3, 2],
15], dtype=tf.int32)
16
17loss = train_step(model, optimizer, token_batch, tag_batch)
18decoded_tags, _ = decode_tags(model, token_batch)
19
20print("loss:", float(loss))
21print(decoded_tags.numpy())

For a real project, replace the dummy tensors with tokenized inputs, gold label IDs, and a proper training loop.

Common Pitfalls

The biggest issue is mishandling sequence lengths. If padded positions are counted as real tokens, the CRF will learn on garbage transitions and both loss and decoding quality will suffer.

Another common mistake is using plain argmax at inference time. That ignores the CRF transition matrix and defeats the main reason for adding a CRF in the first place.

Masking can also be subtle. mask_zero=True helps the embedding layer, but the CRF functions still need explicit sequence lengths.

Finally, check your dependency strategy. TensorFlow Addons still exposes CRF APIs, but it is a separate package from core TensorFlow, so compatibility should be verified in your environment before you build around it.

Summary

• A CRF is useful when label transitions matter across a sequence.
• In TensorFlow 2, a common pattern is encoder logits plus tfa.text.crf_log_likelihood and tfa.text.crf_decode.
• The model predicts emission scores, while the CRF learns transition scores between labels.
• Correct sequence lengths are essential for padded batches.
• Use Viterbi decoding at inference time or you lose the structured prediction benefit of the CRF.