Repository: keras-team/keras-nlp

Issue:
Add an `activation` argument for all classifiers
For general usability and consistency with keras-cv, we would like to add an `activation` argument to all classifiers.

```python
model = BertClassifier.from_preset(
    "bert_base_en_uncased",
    num_classes=2,
    activation="softmax",
)
model.compile(...)
model.fit(...)
model.predict(["What an amazing movie!", "A total waste of my time."])
>>> np.array([[0.004173, 0.9956  ], [0.997   , 0.0028  ]])
```

One way this could work:
- `BertClassifier` activation defaults to `None` (logits).
- If activation is softmax, we default compile with `loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False)`.
- If activation is `None`, we default compile with `loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True)`.
- If activation is something else, we do not default compile.

Discussion:
mattdangerw: /gcbrun
mattdangerw: /gcbrun

Patch:
diff --git a/README.md b/README.md
index 5e16332a..71836262 100644
--- a/README.md
+++ b/README.md
@@ -76,6 +76,7 @@ imdb_train, imdb_test = tfds.load(
 classifier = keras_nlp.models.BertClassifier.from_preset(
     "bert_base_en_uncased", 
     num_classes=2,
+    activation="softmax",
 )
 # Fine-tune on IMDb movie reviews.
 classifier.fit(imdb_train, validation_data=imdb_test)
diff --git a/keras_nlp/models/albert/albert_classifier.py b/keras_nlp/models/albert/albert_classifier.py
index ac708ba4..552e1fc5 100644
--- a/keras_nlp/models/albert/albert_classifier.py
+++ b/keras_nlp/models/albert/albert_classifier.py
@@ -47,11 +47,14 @@ class AlbertClassifier(Task):
     Args:
         backbone: A `keras_nlp.models.AlertBackbone` instance.
         num_classes: int. Number of classes to predict.
-        dropout: float. The dropout probability value, applied after the dense
-            layer.
         preprocessor: A `keras_nlp.models.AlbertPreprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
             be preprocessed before calling the model.
+        activation: Optional `str` or callable, defaults to `None`. The
+            activation function to use on the model outputs. Set
+            `activation="softmax"` to return output probabilities.
+        dropout: float. The dropout probability value, applied after the dense
+            layer.
 
     Examples:
 
@@ -150,8 +153,9 @@ class AlbertClassifier(Task):
         self,
         backbone,
         num_classes,
-        dropout=0.1,
         preprocessor=None,
+        activation=None,
+        dropout=0.1,
         **kwargs,
     ):
         inputs = backbone.input
@@ -160,6 +164,7 @@ class AlbertClassifier(Task):
         outputs = keras.layers.Dense(
             num_classes,
             kernel_initializer=albert_kernel_initializer(),
+            activation=activation,
             name="logits",
         )(pooled)
         # Instantiate using Functional API Model constructor
@@ -173,11 +178,14 @@ class AlbertClassifier(Task):
         self._backbone = backbone
         self._preprocessor = preprocessor
         self.num_classes = num_classes
+        self.activation = activation
         self.dropout = dropout
 
         # Default compilation
         self.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            loss=keras.losses.SparseCategoricalCrossentropy(
+                from_logits=activation is None
+            ),
             optimizer=keras.optimizers.Adam(5e-5),
             metrics=keras.metrics.SparseCategoricalAccuracy(),
             jit_compile=is_xla_compatible(self),
@@ -188,6 +196,7 @@ class AlbertClassifier(Task):
         config.update(
             {
                 "num_classes": self.num_classes,
+                "activation": self.activation,
                 "dropout": self.dropout,
             }
         )
diff --git a/keras_nlp/models/bert/bert_classifier.py b/keras_nlp/models/bert/bert_classifier.py
index 5aac313b..51593ebe 100644
--- a/keras_nlp/models/bert/bert_classifier.py
+++ b/keras_nlp/models/bert/bert_classifier.py
@@ -48,11 +48,14 @@ class BertClassifier(Task):
     Args:
         backbone: A `keras_nlp.models.BertBackbone` instance.
         num_classes: int. Number of classes to predict.
-        dropout: float. The dropout probability value, applied after the dense
-            layer.
         preprocessor: A `keras_nlp.models.BertPreprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
             be preprocessed before calling the model.
+        activation: Optional `str` or callable, defaults to `None`. The
+            activation function to use on the model outputs. Set
+            `activation="softmax"` to return output probabilities.
+        dropout: float. The dropout probability value, applied after the dense
+            layer.
 
     Examples:
 
@@ -138,8 +141,9 @@ class BertClassifier(Task):
         self,
         backbone,
         num_classes,
-        dropout=0.1,
         preprocessor=None,
+        activation=None,
+        dropout=0.1,
         **kwargs,
     ):
         inputs = backbone.input
@@ -148,6 +152,7 @@ class BertClassifier(Task):
         outputs = keras.layers.Dense(
             num_classes,
             kernel_initializer=bert_kernel_initializer(),
+            activation=activation,
             name="logits",
         )(pooled)
         # Instantiate using Functional API Model constructor
@@ -161,11 +166,14 @@ class BertClassifier(Task):
         self.backbone = backbone
         self.preprocessor = preprocessor
         self.num_classes = num_classes
+        self.activation = activation
         self.dropout = dropout
 
         # Default compilation
         self.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            loss=keras.losses.SparseCategoricalCrossentropy(
+                from_logits=activation is None
+            ),
             optimizer=keras.optimizers.Adam(5e-5),
             metrics=keras.metrics.SparseCategoricalAccuracy(),
             jit_compile=is_xla_compatible(self),
@@ -176,6 +184,7 @@ class BertClassifier(Task):
         config.update(
             {
                 "num_classes": self.num_classes,
+                "activation": self.activation,
                 "dropout": self.dropout,
             }
         )
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_classifier.py b/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
index 634ddbd6..7283a4c0 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
@@ -56,12 +56,15 @@ class DebertaV3Classifier(Task):
     Args:
         backbone: A `keras_nlp.models.DebertaV3` instance.
         num_classes: int. Number of classes to predict.
-        hidden_dim: int. The size of the pooler layer.
-        dropout: float. Dropout probability applied to the pooled output. For
-            the second dropout layer, `backbone.dropout` is used.
         preprocessor: A `keras_nlp.models.DebertaV3Preprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
             be preprocessed before calling the model.
+        activation: Optional `str` or callable, defaults to `None`. The
+            activation function to use on the model outputs. Set
+            `activation="softmax"` to return output probabilities.
+        hidden_dim: int. The size of the pooler layer.
+        dropout: float. Dropout probability applied to the pooled output. For
+            the second dropout layer, `backbone.dropout` is used.
 
     Examples:
 
@@ -167,14 +170,14 @@ class DebertaV3Classifier(Task):
         self,
         backbone,
         num_classes,
+        preprocessor=None,
+        activation=None,
         hidden_dim=None,
         dropout=0.0,
-        preprocessor=None,
         **kwargs,
     ):
         inputs = backbone.input
-        if hidden_dim is None:
-            hidden_dim = backbone.hidden_dim
+        hidden_dim = hidden_dim or backbone.hidden_dim
 
         x = backbone(inputs)[:, backbone.start_token_index, :]
         x = keras.layers.Dropout(dropout, name="pooled_dropout")(x)
@@ -187,6 +190,7 @@ class DebertaV3Classifier(Task):
         outputs = keras.layers.Dense(
             num_classes,
             kernel_initializer=deberta_kernel_initializer(),
+            activation=activation,
             name="logits",
         )(x)
 
@@ -201,12 +205,15 @@ class DebertaV3Classifier(Task):
         self.backbone = backbone
         self.preprocessor = preprocessor
         self.num_classes = num_classes
+        self.activation = activation
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
         # Default compilation
         self.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            loss=keras.losses.SparseCategoricalCrossentropy(
+                from_logits=activation is None
+            ),
             optimizer=keras.optimizers.Adam(5e-5),
             metrics=keras.metrics.SparseCategoricalAccuracy(),
             jit_compile=is_xla_compatible(self),
@@ -217,6 +224,7 @@ class DebertaV3Classifier(Task):
         config.update(
             {
                 "num_classes": self.num_classes,
+                "activation": self.activation,
                 "hidden_dim": self.hidden_dim,
                 "dropout": self.dropout,
             }
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier.py b/keras_nlp/models/distil_bert/distil_bert_classifier.py
index 538abf77..f4aab927 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier.py
@@ -53,12 +53,15 @@ class DistilBertClassifier(Task):
     Args:
         backbone: A `keras_nlp.models.DistilBert` instance.
         num_classes: int. Number of classes to predict.
-        hidden_dim: int. The size of the pooler layer.
-        dropout: float. The dropout probability value, applied after the first
-            dense layer.
         preprocessor: A `keras_nlp.models.DistilBertPreprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
             be preprocessed before calling the model.
+        activation: Optional `str` or callable, defaults to `None`. The
+            activation function to use on the model outputs. Set
+            `activation="softmax"` to return output probabilities.
+        hidden_dim: int. The size of the pooler layer.
+        dropout: float. The dropout probability value, applied after the first
+            dense layer.
 
     Examples:
 
@@ -143,15 +146,15 @@ class DistilBertClassifier(Task):
     def __init__(
         self,
         backbone,
-        num_classes=2,
+        num_classes,
+        preprocessor=None,
+        activation=None,
         hidden_dim=None,
         dropout=0.2,
-        preprocessor=None,
         **kwargs,
     ):
         inputs = backbone.input
-        if hidden_dim is None:
-            hidden_dim = backbone.hidden_dim
+        hidden_dim = hidden_dim or backbone.hidden_dim
 
         x = backbone(inputs)[:, backbone.cls_token_index, :]
         x = keras.layers.Dense(
@@ -164,6 +167,7 @@ class DistilBertClassifier(Task):
         outputs = keras.layers.Dense(
             num_classes,
             kernel_initializer=distilbert_kernel_initializer(),
+            activation=activation,
             name="logits",
         )(x)
 
@@ -178,11 +182,14 @@ class DistilBertClassifier(Task):
         self.backbone = backbone
         self.preprocessor = preprocessor
         self.num_classes = num_classes
+        self.activation = activation
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
         self.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            loss=keras.losses.SparseCategoricalCrossentropy(
+                from_logits=activation is None
+            ),
             optimizer=keras.optimizers.Adam(5e-5),
             metrics=keras.metrics.SparseCategoricalAccuracy(),
             jit_compile=is_xla_compatible(self),
@@ -193,6 +200,7 @@ class DistilBertClassifier(Task):
         config.update(
             {
                 "num_classes": self.num_classes,
+                "activation": self.activation,
                 "hidden_dim": self.hidden_dim,
                 "dropout": self.dropout,
             }
diff --git a/keras_nlp/models/f_net/f_net_classifier.py b/keras_nlp/models/f_net/f_net_classifier.py
index 51e23f04..d9c6825d 100644
--- a/keras_nlp/models/f_net/f_net_classifier.py
+++ b/keras_nlp/models/f_net/f_net_classifier.py
@@ -48,12 +48,15 @@ class FNetClassifier(Task):
     Args:
         backbone: A `keras_nlp.models.FNetBackbone` instance.
         num_classes: int. Number of classes to predict.
-        hidden_dim: int. The size of the pooler layer.
-        dropout: float. The dropout probability value, applied after the dense
-            layer.
         preprocessor: A `keras_nlp.models.FNetPreprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
             be preprocessed before calling the model.
+        activation: Optional `str` or callable, defaults to `None`. The
+            activation function to use on the model outputs. Set
+            `activation="softmax"` to return output probabilities.
+        hidden_dim: int. The size of the pooler layer.
+        dropout: float. The dropout probability value, applied after the dense
+            layer.
 
     Examples:
 
@@ -106,8 +109,9 @@ class FNetClassifier(Task):
         self,
         backbone,
         num_classes,
-        dropout=0.1,
         preprocessor=None,
+        activation=None,
+        dropout=0.1,
         **kwargs,
     ):
         inputs = backbone.input
@@ -116,6 +120,7 @@ class FNetClassifier(Task):
         outputs = keras.layers.Dense(
             num_classes,
             kernel_initializer=f_net_kernel_initializer(),
+            activation=activation,
             name="logits",
         )(pooled)
         # Instantiate using Functional API Model constructor
@@ -129,10 +134,13 @@ class FNetClassifier(Task):
         self.backbone = backbone
         self.preprocessor = preprocessor
         self.num_classes = num_classes
+        self.activation = activation
         self.dropout = dropout
 
         self.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            loss=keras.losses.SparseCategoricalCrossentropy(
+                from_logits=activation is None
+            ),
             optimizer=keras.optimizers.Adam(5e-5),
             metrics=keras.metrics.SparseCategoricalAccuracy(),
             jit_compile=is_xla_compatible(self),
@@ -144,6 +152,7 @@ class FNetClassifier(Task):
             {
                 "num_classes": self.num_classes,
                 "dropout": self.dropout,
+                "activation": self.activation,
             }
         )
         return config
diff --git a/keras_nlp/models/roberta/roberta_classifier.py b/keras_nlp/models/roberta/roberta_classifier.py
index 0f711f40..8830e370 100644
--- a/keras_nlp/models/roberta/roberta_classifier.py
+++ b/keras_nlp/models/roberta/roberta_classifier.py
@@ -49,12 +49,15 @@ class RobertaClassifier(Task):
     Args:
         backbone: A `keras_nlp.models.RobertaBackbone` instance.
         num_classes: int. Number of classes to predict.
-        hidden_dim: int. The size of the pooler layer.
-        dropout: float. The dropout probability value, applied to the pooled
-            output, and after the first dense layer.
         preprocessor: A `keras_nlp.models.RobertaPreprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
             be preprocessed before calling the model.
+        activation: Optional `str` or callable, defaults to `None`. The
+            activation function to use on the model outputs. Set
+            `activation="softmax"` to return output probabilities.
+        hidden_dim: int. The size of the pooler layer.
+        dropout: float. The dropout probability value, applied to the pooled
+            output, and after the first dense layer.
 
     Examples:
 
@@ -140,14 +143,14 @@ class RobertaClassifier(Task):
         self,
         backbone,
         num_classes,
+        preprocessor=None,
+        activation=None,
         hidden_dim=None,
         dropout=0.0,
-        preprocessor=None,
         **kwargs,
     ):
         inputs = backbone.input
-        if hidden_dim is None:
-            hidden_dim = backbone.hidden_dim
+        hidden_dim = hidden_dim or backbone.hidden_dim
 
         x = backbone(inputs)[:, backbone.start_token_index, :]
         x = keras.layers.Dropout(dropout, name="pooled_dropout")(x)
@@ -158,6 +161,7 @@ class RobertaClassifier(Task):
         outputs = keras.layers.Dense(
             num_classes,
             kernel_initializer=roberta_kernel_initializer(),
+            activation=activation,
             name="logits",
         )(x)
 
@@ -172,12 +176,15 @@ class RobertaClassifier(Task):
         self.backbone = backbone
         self.preprocessor = preprocessor
         self.num_classes = num_classes
+        self.activation = activation
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
         # Default compilation
         self.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            loss=keras.losses.SparseCategoricalCrossentropy(
+                from_logits=activation is None
+            ),
             optimizer=keras.optimizers.Adam(2e-5),
             metrics=keras.metrics.SparseCategoricalAccuracy(),
             jit_compile=is_xla_compatible(self),
@@ -188,6 +195,7 @@ class RobertaClassifier(Task):
         config.update(
             {
                 "num_classes": self.num_classes,
+                "activation": self.activation,
                 "hidden_dim": self.hidden_dim,
                 "dropout": self.dropout,
             }
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
index abc570d6..d8481b07 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
@@ -51,12 +51,15 @@ class XLMRobertaClassifier(Task):
     Args:
         backbone: A `keras_nlp.models.XLMRobertaBackbone` instance.
         num_classes: int. Number of classes to predict.
-        hidden_dim: int. The size of the pooler layer.
-        dropout: float. The dropout probability value, applied to the pooled
-            output, and after the first dense layer.
         preprocessor: A `keras_nlp.models.XLMRobertaPreprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
             be preprocessed before calling the model.
+        activation: Optional `str` or callable, defaults to `None`. The
+            activation function to use on the model outputs. Set
+            `activation="softmax"` to return output probabilities.
+        hidden_dim: int. The size of the pooler layer.
+        dropout: float. The dropout probability value, applied to the pooled
+            output, and after the first dense layer.
 
     Examples:
 
@@ -153,14 +156,14 @@ class XLMRobertaClassifier(Task):
         self,
         backbone,
         num_classes,
+        preprocessor=None,
+        activation=None,
         hidden_dim=None,
         dropout=0.0,
-        preprocessor=None,
         **kwargs,
     ):
         inputs = backbone.input
-        if hidden_dim is None:
-            hidden_dim = backbone.hidden_dim
+        hidden_dim = hidden_dim or backbone.hidden_dim
 
         x = backbone(inputs)[:, backbone.start_token_index, :]
         x = keras.layers.Dropout(dropout, name="pooled_dropout")(x)
@@ -171,6 +174,7 @@ class XLMRobertaClassifier(Task):
         outputs = keras.layers.Dense(
             num_classes,
             kernel_initializer=roberta_kernel_initializer(),
+            activation=activation,
             name="logits",
         )(x)
 
@@ -185,11 +189,14 @@ class XLMRobertaClassifier(Task):
         self.backbone = backbone
         self.preprocessor = preprocessor
         self.num_classes = num_classes
+        self.activation = activation
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
         self.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            loss=keras.losses.SparseCategoricalCrossentropy(
+                from_logits=activation is None
+            ),
             optimizer=keras.optimizers.Adam(5e-5),
             metrics=keras.metrics.SparseCategoricalAccuracy(),
             jit_compile=is_xla_compatible(self),
@@ -203,6 +210,7 @@ class XLMRobertaClassifier(Task):
         config.update(
             {
                 "num_classes": self.num_classes,
+                "activation": self.activation,
                 "hidden_dim": self.hidden_dim,
                 "dropout": self.dropout,
             }


Test patch:
diff --git a/keras_nlp/models/albert/albert_classifier_test.py b/keras_nlp/models/albert/albert_classifier_test.py
index af84805b..bbf34d02 100644
--- a/keras_nlp/models/albert/albert_classifier_test.py
+++ b/keras_nlp/models/albert/albert_classifier_test.py
@@ -71,26 +71,20 @@ class AlbertClassifierTest(tf.test.TestCase, parameterized.TestCase):
 
         self.classifier = AlbertClassifier(
             self.backbone,
-            4,
+            num_classes=4,
             preprocessor=self.preprocessor,
-        )
-        self.classifier_no_preprocessing = AlbertClassifier(
-            self.backbone,
-            4,
-            preprocessor=None,
+            activation="softmax",
         )
 
         self.raw_batch = tf.constant(
             [
                 "the quick brown fox.",
                 "the slow brown fox.",
-                "the smelly brown fox.",
-                "the old brown fox.",
             ]
         )
         self.preprocessed_batch = self.preprocessor(self.raw_batch)
         self.raw_dataset = tf.data.Dataset.from_tensor_slices(
-            (self.raw_batch, tf.ones((4,)))
+            (self.raw_batch, tf.ones((2,)))
         ).batch(2)
         self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
 
@@ -98,17 +92,26 @@ class AlbertClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier(self.preprocessed_batch)
 
     def test_classifier_predict(self):
-        self.classifier.predict(self.raw_batch)
+        preds1 = self.classifier.predict(self.raw_batch)
         self.classifier.preprocessor = None
-        self.classifier.predict(self.preprocessed_batch)
+        preds2 = self.classifier.predict(self.preprocessed_batch)
+        # Assert predictions match.
+        self.assertAllClose(preds1, preds2)
+        # Assert valid softmax output.
+        self.assertAllClose(tf.reduce_sum(preds2, axis=-1), [1.0, 1.0])
 
     def test_classifier_fit(self):
         self.classifier.fit(self.raw_dataset)
         self.classifier.preprocessor = None
         self.classifier.fit(self.preprocessed_dataset)
 
-    def test_distilbert_classifier_fit_default_compile(self):
-        self.classifier.fit(self.raw_dataset)
+    def test_classifier_fit_no_xla(self):
+        self.classifier.preprocessor = None
+        self.classifier.compile(
+            loss="sparse_categorical_crossentropy",
+            jit_compile=False,
+        )
+        self.classifier.fit(self.preprocessed_dataset)
 
     def test_serialization(self):
         config = keras.utils.serialize_keras_object(self.classifier)
diff --git a/keras_nlp/models/bert/bert_classifier_test.py b/keras_nlp/models/bert/bert_classifier_test.py
index 11c25431..cda85bc5 100644
--- a/keras_nlp/models/bert/bert_classifier_test.py
+++ b/keras_nlp/models/bert/bert_classifier_test.py
@@ -45,8 +45,9 @@ class BertClassifierTest(tf.test.TestCase, parameterized.TestCase):
         )
         self.classifier = BertClassifier(
             self.backbone,
-            4,
+            num_classes=4,
             preprocessor=self.preprocessor,
+            activation="softmax",
         )
 
         # Setup data.
@@ -66,9 +67,13 @@ class BertClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier(self.preprocessed_batch)
 
     def test_classifier_predict(self):
-        self.classifier.predict(self.raw_batch)
+        preds1 = self.classifier.predict(self.raw_batch)
         self.classifier.preprocessor = None
-        self.classifier.predict(self.preprocessed_batch)
+        preds2 = self.classifier.predict(self.preprocessed_batch)
+        # Assert predictions match.
+        self.assertAllClose(preds1, preds2)
+        # Assert valid softmax output.
+        self.assertAllClose(tf.reduce_sum(preds2, axis=-1), [1.0, 1.0])
 
     def test_classifier_fit(self):
         self.classifier.fit(self.raw_dataset)
@@ -78,7 +83,7 @@ class BertClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def test_classifier_fit_no_xla(self):
         self.classifier.preprocessor = None
         self.classifier.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False),
+            loss="sparse_categorical_crossentropy",
             jit_compile=False,
         )
         self.classifier.fit(self.preprocessed_dataset)
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py b/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
index 9c255d21..72c4d65c 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
@@ -68,8 +68,10 @@ class DebertaV3ClassifierTest(tf.test.TestCase, parameterized.TestCase):
         )
         self.classifier = DebertaV3Classifier(
             self.backbone,
-            4,
+            num_classes=4,
             preprocessor=self.preprocessor,
+            activation="softmax",
+            hidden_dim=4,
         )
 
         self.raw_batch = tf.constant(
@@ -88,9 +90,13 @@ class DebertaV3ClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier(self.preprocessed_batch)
 
     def test_classifier_predict(self):
-        self.classifier.predict(self.raw_batch)
+        preds1 = self.classifier.predict(self.raw_batch)
         self.classifier.preprocessor = None
-        self.classifier.predict(self.preprocessed_batch)
+        preds2 = self.classifier.predict(self.preprocessed_batch)
+        # Assert predictions match.
+        self.assertAllClose(preds1, preds2)
+        # Assert valid softmax output.
+        self.assertAllClose(tf.reduce_sum(preds2, axis=-1), [1.0, 1.0])
 
     def test_classifier_fit(self):
         self.classifier.fit(self.raw_dataset)
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
index 9e1eda99..e39b149d 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
@@ -52,21 +52,21 @@ class DistilBertClassifierTest(tf.test.TestCase, parameterized.TestCase):
         )
         self.classifier = DistilBertClassifier(
             self.backbone,
-            4,
+            num_classes=4,
             preprocessor=self.preprocessor,
+            activation="softmax",
+            hidden_dim=4,
         )
 
         self.raw_batch = tf.constant(
             [
                 "the quick brown fox.",
                 "the slow brown fox.",
-                "the smelly brown fox.",
-                "the old brown fox.",
             ]
         )
         self.preprocessed_batch = self.preprocessor(self.raw_batch)
         self.raw_dataset = tf.data.Dataset.from_tensor_slices(
-            (self.raw_batch, tf.ones((4,)))
+            (self.raw_batch, tf.ones((2,)))
         ).batch(2)
         self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
 
@@ -74,17 +74,26 @@ class DistilBertClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier(self.preprocessed_batch)
 
     def test_classifier_predict(self):
-        self.classifier.predict(self.raw_batch)
+        preds1 = self.classifier.predict(self.raw_batch)
         self.classifier.preprocessor = None
-        self.classifier.predict(self.preprocessed_batch)
+        preds2 = self.classifier.predict(self.preprocessed_batch)
+        # Assert predictions match.
+        self.assertAllClose(preds1, preds2)
+        # Assert valid softmax output.
+        self.assertAllClose(tf.reduce_sum(preds2, axis=-1), [1.0, 1.0])
 
     def test_classifier_fit(self):
         self.classifier.fit(self.raw_dataset)
         self.classifier.preprocessor = None
         self.classifier.fit(self.preprocessed_dataset)
 
-    def test_distilbert_classifier_fit_default_compile(self):
-        self.classifier.fit(self.raw_dataset)
+    def test_classifier_fit_no_xla(self):
+        self.classifier.preprocessor = None
+        self.classifier.compile(
+            loss="sparse_categorical_crossentropy",
+            jit_compile=False,
+        )
+        self.classifier.fit(self.preprocessed_dataset)
 
     def test_serialization(self):
         config = keras.utils.serialize_keras_object(self.classifier)
diff --git a/keras_nlp/models/distil_bert/distil_bert_presets_test.py b/keras_nlp/models/distil_bert/distil_bert_presets_test.py
index 8b151893..f3a25ad6 100644
--- a/keras_nlp/models/distil_bert/distil_bert_presets_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_presets_test.py
@@ -78,6 +78,7 @@ class DistilBertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
         input_data = tf.constant(["The quick brown fox."])
         model = DistilBertClassifier.from_preset(
             "distil_bert_base_en_uncased",
+            num_classes=2,
             load_weights=load_weights,
         )
         model.predict(input_data)
@@ -92,6 +93,7 @@ class DistilBertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
         }
         model = DistilBertClassifier.from_preset(
             "distil_bert_base_en_uncased",
+            num_classes=2,
             load_weights=load_weights,
             preprocessor=None,
         )
@@ -153,7 +155,7 @@ class DistilBertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
         for preset in DistilBertClassifier.presets:
             classifier = DistilBertClassifier.from_preset(
                 preset,
-                num_classes=4,
+                num_classes=2,
                 load_weights=load_weights,
             )
             input_data = tf.constant(["This quick brown fox"])
@@ -166,7 +168,7 @@ class DistilBertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
         for preset in DistilBertClassifier.presets:
             classifier = DistilBertClassifier.from_preset(
                 preset,
-                num_classes=4,
+                num_classes=2,
                 load_weights=load_weights,
                 preprocessor=None,
             )
diff --git a/keras_nlp/models/f_net/f_net_classifier_test.py b/keras_nlp/models/f_net/f_net_classifier_test.py
index d6dbaf36..f1b179c3 100644
--- a/keras_nlp/models/f_net/f_net_classifier_test.py
+++ b/keras_nlp/models/f_net/f_net_classifier_test.py
@@ -67,8 +67,9 @@ class FNetClassifierTest(tf.test.TestCase, parameterized.TestCase):
         )
         self.classifier = FNetClassifier(
             self.backbone,
-            4,
+            num_classes=4,
             preprocessor=self.preprocessor,
+            activation="softmax",
         )
 
         # Setup data.
@@ -88,9 +89,13 @@ class FNetClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier(self.preprocessed_batch)
 
     def test_classifier_predict(self):
-        self.classifier.predict(self.raw_batch)
+        preds1 = self.classifier.predict(self.raw_batch)
         self.classifier.preprocessor = None
-        self.classifier.predict(self.preprocessed_batch)
+        preds2 = self.classifier.predict(self.preprocessed_batch)
+        # Assert predictions match.
+        self.assertAllClose(preds1, preds2)
+        # Assert valid softmax output.
+        self.assertAllClose(tf.reduce_sum(preds2, axis=-1), [1.0, 1.0])
 
     def test_fnet_classifier_fit(self):
         self.classifier.fit(self.raw_dataset)
@@ -120,8 +125,8 @@ class FNetClassifierTest(tf.test.TestCase, parameterized.TestCase):
     @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
-        path = os.path.join(self.get_temp_dir(), filename)
-        # Don't save traces in the tf format, we check compilation elsewhere.
+        path = os.path.join(self.get_temp_dir(), filename)
+        # Don't save traces in the tf format, we check compilation elsewhere.
         kwargs = {"save_traces": False} if save_format == "tf" else {}
         self.classifier.save(path, save_format=save_format, **kwargs)
         restored_model = keras.models.load_model(path)
diff --git a/keras_nlp/models/roberta/roberta_classifier_test.py b/keras_nlp/models/roberta/roberta_classifier_test.py
index bc1b0ca9..f2eb012e 100644
--- a/keras_nlp/models/roberta/roberta_classifier_test.py
+++ b/keras_nlp/models/roberta/roberta_classifier_test.py
@@ -63,8 +63,10 @@ class RobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
         )
         self.classifier = RobertaClassifier(
             self.backbone,
-            4,
+            num_classes=4,
             preprocessor=self.preprocessor,
+            activation="softmax",
+            hidden_dim=4,
         )
 
         # Setup data.
@@ -84,9 +86,13 @@ class RobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier(self.preprocessed_batch)
 
     def test_classifier_predict(self):
-        self.classifier.predict(self.raw_batch)
+        preds1 = self.classifier.predict(self.raw_batch)
         self.classifier.preprocessor = None
-        self.classifier.predict(self.preprocessed_batch)
+        preds2 = self.classifier.predict(self.preprocessed_batch)
+        # Assert predictions match.
+        self.assertAllClose(preds1, preds2)
+        # Assert valid softmax output.
+        self.assertAllClose(tf.reduce_sum(preds2, axis=-1), [1.0, 1.0])
 
     def test_classifier_fit(self):
         self.classifier.fit(self.raw_dataset)
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
index 082b2eea..a96fa8d9 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
@@ -63,13 +63,10 @@ class XLMRobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
         )
         self.classifier = XLMRobertaClassifier(
             self.backbone,
-            4,
+            num_classes=4,
             preprocessor=self.preprocessor,
-        )
-        self.classifier_no_preprocessing = XLMRobertaClassifier(
-            self.backbone,
-            4,
-            preprocessor=None,
+            activation="softmax",
+            hidden_dim=4,
         )
 
         self.raw_batch = tf.constant(
@@ -88,9 +85,13 @@ class XLMRobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier(self.preprocessed_batch)
 
     def test_classifier_predict(self):
-        self.classifier.predict(self.raw_batch)
+        preds1 = self.classifier.predict(self.raw_batch)
         self.classifier.preprocessor = None
-        self.classifier.predict(self.preprocessed_batch)
+        preds2 = self.classifier.predict(self.preprocessed_batch)
+        # Assert predictions match.
+        self.assertAllClose(preds1, preds2)
+        # Assert valid softmax output.
+        self.assertAllClose(tf.reduce_sum(preds2, axis=-1), [1.0, 1.0])
 
     def test_classifier_fit(self):
         self.classifier.fit(self.raw_dataset)
@@ -100,7 +101,7 @@ class XLMRobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def test_classifier_fit_no_xla(self):
         self.classifier.preprocessor = None
         self.classifier.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False),
+            loss="sparse_categorical_crossentropy",
             jit_compile=False,
         )
         self.classifier.fit(self.preprocessed_dataset)