Introduction

In TensorFlow 2, a CRF is usually not added as a normal Dense(..., activation=...) layer. Instead, you produce per-token logits with your network, then use tfa.text.crf_log_likelihood for training and tfa.text.crf_decode for inference.

Understand What the CRF Needs

A CRF layer is useful for sequence labeling tasks such as named entity recognition because it models dependencies between adjacent output tags. The network produces emission scores for each token, and the CRF adds transition scores between labels.

That means the model needs three core pieces:

• logits shaped like [batch, time, num_tags]
• true tag ids shaped like [batch, time]
• sequence lengths so padding is ignored correctly

The transition matrix is typically a trainable variable of shape [num_tags, num_tags].

Install Matching Packages

TensorFlow Addons provides the CRF utilities under tfa.text. The package version must match the TensorFlow version closely enough to be supported.

pip install tensorflow tensorflow-addons

If versions are mismatched, CRF code can fail in ways that look unrelated to your model. Check compatibility first when debugging import or runtime issues.

Build Logits with a Keras Model

The CRF itself is easiest to manage when the neural network produces logits and a custom training step handles the CRF loss.

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

This model stops at logits. That is intentional because CRF training uses a special loss, not a normal softmax loss.

Compute CRF Log Likelihood

During training, compute sequence lengths from non-padding tokens and pass them to crf_log_likelihood.

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

The negative mean log likelihood is the training loss. Padding must be handled correctly or the CRF will learn transitions across fake tokens.

Decode with Viterbi

At inference time, use crf_decode instead of taking argmax over logits independently.

1def predict_tags(model, token_ids):
2    sequence_lengths = tf.reduce_sum(tf.cast(token_ids != 0, tf.int32), axis=1)
3    logits = model(token_ids, training=False)
4    decoded_tags, scores = tfa.text.crf_decode(
5        logits,
6        model.transition_params,
7        sequence_lengths,
8    )
9    return decoded_tags, scores

This is the whole point of the CRF: decoding considers the best sequence globally, not each label position in isolation.

Practical Notes

Sequence masking is the main operational detail. If your padding token is not 0, change the sequence-length calculation accordingly. Also note that TensorFlow Addons has been in maintenance mode, so teams starting greenfield work should check project support status before depending on it heavily.

That does not make tfa.text unusable, but it does mean you should be deliberate about version pinning.

Common Pitfalls

• Treating the CRF like a normal activation layer and trying to train it with plain categorical cross-entropy.
• Forgetting to pass correct sequence lengths, which causes padding tokens to distort the loss and decode path.
• Using argmax at inference time instead of crf_decode.
• Leaving the transition matrix out of the trainable variables used by the optimizer.
• Ignoring TensorFlow and TensorFlow Addons version compatibility when setup errors appear.

Summary

• In TensorFlow 2, a CRF is typically implemented with CRF utilities around network logits, not as a simple final activation.
• Use tfa.text.crf_log_likelihood for training and tfa.text.crf_decode for inference.
• Sequence lengths are essential so padded tokens are excluded correctly.
• Keep the CRF transition matrix trainable and included in optimization.
• Version compatibility and project support status matter when using TensorFlow Addons.