Repository: keras-team/keras-nlp
Pull request #1201 (2023-08-09 00:25:15)

Add support for "untied" embedding weights in language models
This PR adds support for language models with "untied weights"--where the weights for the output projection are separate from the weights for the input embedding. We need this for T5 and some other newer models.

To do this, we will extend our token embedding to also support a "reverse projection." We will introduce a new `ReversibleEmbedding` layer, and if `reverse=True` is passed to call, we project from `output_dim` back to `input_dim`. Handling this all in the token embedding layer allows us to have an easy setup where one layer clearly owns all the large embedding matrices.

We also need to introduce a backwards incompatible change for the `MaskedLMHead` layer, to take as input a layer instead of a shared variable. From discussion on https://github.com/keras-team/keras-core/issues/681, passing a shared variable to init is not actually a good approach; we can rely on layers being deduplicated properly, but not weights. All our masked lm tasks will work unaltered (e.g. `keras_nlp.models.BertMaskedLM`), this API change is just for the layer.

Side note, this gives us a nice recipe for converting a `backbone` to a simple language model.
```python
backbone = ...
inputs = backbone.input
hidden_states = backbone(inputs)
output = backbone.token_embedding(hidden_states, reverse=True)
language_model = keras.Model(inputs, outputs)
```

mattdangerw: /gcbrun
mattdangerw: /gcbrun
mattdangerw: /gcbrun
mattdangerw: /gcbrun
mattdangerw: /gcbrun
abheesht17: Name's not `TransposedEmbedding`? :(

(just kidding :P)
abheesht17: Name's not `TransposedEmbedding`? :(

(just kidding :P)

Patch:
diff --git a/keras_nlp/layers/modeling/masked_lm_head.py b/keras_nlp/layers/modeling/masked_lm_head.py
index e8a7e4a6..46c15f02 100644
--- a/keras_nlp/layers/modeling/masked_lm_head.py
+++ b/keras_nlp/layers/modeling/masked_lm_head.py
@@ -26,7 +26,7 @@ class MaskedLMHead(keras.layers.Layer):
     This layer takes two inputs:
 
      - `inputs`: which should be a tensor of encoded tokens with shape
-       `(batch_size, sequence_length, encoding_dim)`.
+       `(batch_size, sequence_length, hidden_dim)`.
      - `mask_positions`: which should be a tensor of integer positions to
        predict with shape `(batch_size, masks_per_sequence)`.
 
@@ -46,11 +46,9 @@ class MaskedLMHead(keras.layers.Layer):
 
     Args:
         vocabulary_size: The total size of the vocabulary for predictions.
-        embedding_weights: Optional. The weights of the word embedding used
-            to transform input token ids. The transpose of this weight matrix
-            will be used to project a token embedding vector to a prediction
-            over all input words, as described
-            [here](https://arxiv.org/abs/1608.05859).
+        token_embedding: Optional. A `keras_nlp.layers.ReversibleEmbedding`
+            instance. If passed, the layer will be used to project from the
+            `hidden_dim` of the model to the output `vocabulary_size`.
         intermediate_activation: The activation function of inner dense layer.
         activation: The activation function for the outputs of the layer.
             Usually either `None` (return logits), or `"softmax"`
@@ -63,37 +61,32 @@ class MaskedLMHead(keras.layers.Layer):
         bias_initializer: string or `keras.initializers` initializer.
             The bias initializer for the dense and multiheaded
             attention layers. Defaults to `"zeros"`.
-        name: string. The name of the layer. Defaults to `None`.
-        **kwargs: other keyword arguments.
 
     Examples:
 
     ```python
-    batch_size = 32
+    batch_size = 16
     vocab_size = 100
-    encoding_size = 32
+    hidden_dim = 32
     seq_length = 50
-    mask_length = 10
 
-    # Generate a random encoding.
-    encoded_tokens = np.random.normal(
-        size=(batch_size, seq_length, encoding_size),
-    )
-    # Generate random positions and labels
-    mask_positions = np.random.randint(
-        seq_length, size=(batch_size, mask_length),
-    )
-    mask_ids = np.random.randint(
-        vocab_size, size=(batch_size, mask_length),
+    # Generate random inputs.
+    token_ids = np.random.randint(vocab_size, size=(batch_size, seq_length))
+    # Choose random positions as the masked inputs.
+    mask_positions = np.random.randint(seq_length, size=(batch_size, 5))
+
+    # Embed tokens in a `hidden_dim` feature space.
+    token_embedding = keras_nlp.layers.ReversibleEmbedding(
+        vocab_size,
+        hidden_dim,
     )
+    hidden_states = token_embedding(token_ids)
 
-    # Predict an output word for each masked input token.
-    mask_preds = keras_nlp.layers.MaskedLMHead(
+    preds = keras_nlp.layers.MaskedLMHead(
         vocabulary_size=vocab_size,
+        token_embedding=token_embedding,
         activation="softmax",
-    )(encoded_tokens, mask_positions=mask_positions)
-    # Calculate a loss.
-    keras.losses.sparse_categorical_crossentropy(mask_ids, mask_preds)
+    )(hidden_states, mask_positions)
     ```
 
     References:
@@ -103,19 +96,18 @@ class MaskedLMHead(keras.layers.Layer):
     def __init__(
         self,
         vocabulary_size=None,
-        embedding_weights=None,
+        token_embedding=None,
         intermediate_activation="relu",
         activation=None,
         layer_norm_epsilon=1e-05,
         kernel_initializer="glorot_uniform",
         bias_initializer="zeros",
-        name=None,
         **kwargs,
     ):
-        super().__init__(name=name, **kwargs)
+        super().__init__(**kwargs)
 
         self.vocabulary_size = vocabulary_size
-        self.embedding_weights = embedding_weights
+        self.token_embedding = token_embedding
         self.intermediate_activation = keras.activations.get(
             intermediate_activation
         )
@@ -124,26 +116,25 @@ class MaskedLMHead(keras.layers.Layer):
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
 
-        if vocabulary_size is None and embedding_weights is None:
+        if vocabulary_size is None and token_embedding is None:
             raise ValueError(
-                "One of `vocabulary_size` or `embedding_weights` must be set. "
-                "Received: `vocabulary_size=None`, `embedding_weights=None`"
+                "One of `vocabulary_size` or `token_embedding` must be set. "
+                "Received: `vocabulary_size=None`, `token_embedding=None`"
             )
 
-        if embedding_weights is not None:
-            shape = embedding_weights.shape
-            if vocabulary_size is not None and vocabulary_size != shape[0]:
+        if token_embedding:
+            if vocabulary_size and vocabulary_size != token_embedding.input_dim:
                 raise ValueError(
-                    "`vocabulary_size` should match the first dimension of the "
-                    "shape of `embedding_weights`. Received: "
+                    "`vocabulary_size` should match the input dimension of the "
+                    "of `token_embedding`. Received: "
                     f"`vocabulary_size={vocabulary_size}`, "
-                    f"`embedding_weights.shape={shape}`"
+                    f"`token_embedding.input_dim={token_embedding.input_dim}`"
                 )
-            self.vocabulary_size = shape[0]
+            self.vocabulary_size = token_embedding.input_dim
 
     def build(self, inputs_shape, mask_positions_shape=None):
-        if self.embedding_weights is not None:
-            feature_size = self.embedding_weights.shape[-1]
+        if self.token_embedding is not None:
+            feature_size = self.token_embedding.output_dim
         else:
             feature_size = inputs_shape[-1]
 
@@ -163,7 +154,7 @@ class MaskedLMHead(keras.layers.Layer):
         self._dense.build(shape)
         shape = (inputs_shape[0], gather_length, feature_size)
         self._layer_norm.build(shape)
-        if self.embedding_weights is None:
+        if self.token_embedding is None:
             self._kernel = self.add_weight(
                 name="output_kernel",
                 shape=[feature_size, self.vocabulary_size],
@@ -187,12 +178,10 @@ class MaskedLMHead(keras.layers.Layer):
         x = self._layer_norm(x)
 
         # Transform encodings to vocabulary_size predictions.
-        if self.embedding_weights is None:
-            kernel = self._kernel
+        if self.token_embedding:
+            outputs = self.token_embedding(x, reverse=True)
         else:
-            kernel = ops.cast(self.embedding_weights, self.compute_dtype)
-            kernel = ops.transpose(kernel)
-        outputs = ops.matmul(x, kernel)
+            outputs = ops.matmul(x, self._kernel)
         outputs = outputs + self._bias
 
         # Apply a final activation.
@@ -201,11 +190,22 @@ class MaskedLMHead(keras.layers.Layer):
 
         return outputs
 
+    @classmethod
+    def from_config(cls, config):
+        embedding = config.get("token_embedding")
+        if embedding:
+            config["token_embedding"] = keras.layers.deserialize(embedding)
+        return super().from_config(config)
+
     def get_config(self):
         config = super().get_config()
+        embedding_config = None
+        if self.token_embedding:
+            embedding_config = keras.layers.serialize(self.token_embedding)
         config.update(
             {
                 "vocabulary_size": self.vocabulary_size,
+                "token_embedding": embedding_config,
                 "intermediate_activation": keras.activations.serialize(
                     self.intermediate_activation
                 ),
diff --git a/keras_nlp/layers/modeling/reversible_embedding.py b/keras_nlp/layers/modeling/reversible_embedding.py
new file mode 100644
index 00000000..fd9e9e3a
--- /dev/null
+++ b/keras_nlp/layers/modeling/reversible_embedding.py
@@ -0,0 +1,153 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+
+from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import keras
+from keras_nlp.backend import ops
+
+
+@keras_nlp_export("keras_nlp.layers.ReversibleEmbedding")
+class ReversibleEmbedding(keras.layers.Embedding):
+    """An embedding layer which can project backwards to the input dim.
+
+    This layer is an extension of `keras.layers.Embedding` for language models.
+    This layer can be called "in reverse" with `reverse=True`, in which case the
+    layer will linearly project from `output_dim` back to `input_dim`.
+
+    By default, the reverse projection will use the transpose of the
+    `embeddings` weights to project to `input_dim` (weights are "tied"). If
+    `tie_weights=False`, the model will use a separate, trainable variable for
+    reverse projection.
+
+    This layer has no bias terms.
+
+    Args:
+        input_dim: Integer. Size of the vocabulary,
+            i.e. maximum integer index + 1.
+        output_dim: Integer. Dimension of the dense embedding.
+        tie_weights: Boolean, whether or not the matrix for embedding and
+            the matrix for the `reverse` projection should share the same
+            weights.
+        embeddings_initializer: Initializer for the `embeddings`
+            matrix (see `keras.initializers`).
+        embeddings_regularizer: Regularizer function applied to
+            the `embeddings` matrix (see `keras.regularizers`).
+        embeddings_constraint: Constraint function applied to
+            the `embeddings` matrix (see `keras.constraints`).
+        mask_zero: Boolean, whether or not the input value 0 is a special
+            "padding" value that should be masked out.
+
+    Call args:
+        inputs: The tensor inputs to the layer.
+        reverse: Boolean. If `True` the layer will perform a linear projection
+            from `output_dim` to `input_dim`, instead of a normal embedding
+            call. Default to `False`.
+
+    Examples:
+    ```python
+    batch_size = 16
+    vocab_size = 100
+    hidden_dim = 32
+    seq_length = 50
+
+    # Generate random inputs.
+    token_ids = np.random.randint(vocab_size, size=(batch_size, seq_length))
+
+    embedding = keras.layers.Embedding(vocab_size, hidden_dim)
+    # Embed tokens to shape `(batch_size, seq_length, hidden_dim)`.
+    hidden_states = embedding(token_ids)
+    # Project hidden states to shape `(batch_size, seq_length, vocab_size)`.
+    logits = embedding(hidden_state, reverse=True)
+    ```
+
+    References:
+    - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
+    - [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
+    """
+
+    def __init__(
+        self,
+        input_dim,
+        output_dim,
+        tie_weights=True,
+        embeddings_initializer="uniform",
+        embeddings_regularizer=None,
+        embeddings_constraint=None,
+        mask_zero=False,
+        **kwargs,
+    ):
+        super().__init__(
+            input_dim,
+            output_dim,
+            embeddings_initializer=embeddings_initializer,
+            embeddings_regularizer=embeddings_regularizer,
+            embeddings_constraint=embeddings_constraint,
+            mask_zero=mask_zero,
+            **kwargs,
+        )
+        self.tie_weights = tie_weights
+
+    def build(self, inputs_shape=None):
+        super().build(inputs_shape)
+
+        if not self.tie_weights:
+            self.reverse_embeddings = self.add_weight(
+                name="reverse_embeddings",
+                shape=(self.output_dim, self.input_dim),
+                initializer=self.embeddings_initializer,
+                dtype=self.dtype,
+            )
+
+    def call(self, inputs, reverse=False):
+        if reverse:
+            if self.tie_weights:
+                reverse_embeddings = ops.transpose(
+                    ops.convert_to_tensor(self.embeddings)
+                )
+            else:
+                reverse_embeddings = self.reverse_embeddings
+            return ops.matmul(inputs, reverse_embeddings)
+
+        return super().call(inputs)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "tie_weights": self.tie_weights,
+            }
+        )
+        return config
+
+    def load_own_variables(self, store):
+        if not self.built:
+            self.build()
+        self.embeddings.assign(store["0"])
+        if not self.tie_weights:
+            # Handle the case where saved weights are tied, but the layer
+            # weights untied. We can simply assign the embedding weights to both
+            # variables in this case.
+            if len(store.keys()) == 1:
+                self.reverse_embeddings.assign(np.transpose(store["0"]))
+            else:
+                self.reverse_embeddings.assign(store["1"])
+
+    def compute_output_spec(self, inputs, reverse=False):
+        output_shape = list(inputs.shape)
+        if reverse:
+            output_shape[-1] = self.input_dim
+        else:
+            output_shape += [self.output_dim]
+        return keras.KerasTensor(output_shape, dtype=self.dtype)
diff --git a/keras_nlp/layers/modeling/token_and_position_embedding.py b/keras_nlp/layers/modeling/token_and_position_embedding.py
index 216a42ff..aebf9090 100644
--- a/keras_nlp/layers/modeling/token_and_position_embedding.py
+++ b/keras_nlp/layers/modeling/token_and_position_embedding.py
@@ -17,6 +17,7 @@
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.utils.keras_utils import clone_initializer
 
 
@@ -62,6 +63,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
         vocabulary_size,
         sequence_length,
         embedding_dim,
+        tie_weights=True,
         embeddings_initializer="glorot_uniform",
         mask_zero=False,
         **kwargs
@@ -85,9 +87,10 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
         self.embeddings_initializer = keras.initializers.get(
             embeddings_initializer
         )
-        self.token_embedding = keras.layers.Embedding(
+        self.token_embedding = ReversibleEmbedding(
             vocabulary_size,
             embedding_dim,
+            tie_weights=tie_weights,
             embeddings_initializer=clone_initializer(
                 self.embeddings_initializer
             ),
@@ -119,6 +122,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
                 "embeddings_initializer": keras.initializers.serialize(
                     self.embeddings_initializer
                 ),
+                "tie_weights": self.token_embedding.tie_weights,
                 "mask_zero": self.token_embedding.mask_zero,
             },
         )
diff --git a/keras_nlp/models/albert/albert_backbone.py b/keras_nlp/models/albert/albert_backbone.py
index 36424faf..45da2d30 100644
--- a/keras_nlp/models/albert/albert_backbone.py
+++ b/keras_nlp/models/albert/albert_backbone.py
@@ -19,6 +19,7 @@ import copy
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_nlp.models.albert.albert_presets import backbone_presets
 from keras_nlp.models.backbone import Backbone
@@ -132,7 +133,7 @@ class AlbertBackbone(Backbone):
         )
 
         # Embed tokens, positions, and segment ids.
-        token_embedding_layer = keras.layers.Embedding(
+        token_embedding_layer = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=embedding_dim,
             embeddings_initializer=albert_kernel_initializer(),
@@ -246,6 +247,7 @@ class AlbertBackbone(Backbone):
         self.max_sequence_length = max_sequence_length
         self.num_segments = num_segments
         self.cls_token_index = cls_token_index
+        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
@@ -266,10 +268,6 @@ class AlbertBackbone(Backbone):
         )
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_embedding")
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/albert/albert_masked_lm.py b/keras_nlp/models/albert/albert_masked_lm.py
index fa04b83b..8084340a 100644
--- a/keras_nlp/models/albert/albert_masked_lm.py
+++ b/keras_nlp/models/albert/albert_masked_lm.py
@@ -108,7 +108,7 @@ class AlbertMaskedLM(Task):
         backbone_outputs = backbone(backbone.input)
         outputs = MaskedLMHead(
             vocabulary_size=backbone.vocabulary_size,
-            embedding_weights=backbone.token_embedding.embeddings,
+            token_embedding=backbone.token_embedding,
             intermediate_activation=lambda x: keras.activations.gelu(
                 x, approximate=True
             ),
diff --git a/keras_nlp/models/backbone.py b/keras_nlp/models/backbone.py
index 66c44292..26cd4bf8 100644
--- a/keras_nlp/models/backbone.py
+++ b/keras_nlp/models/backbone.py
@@ -25,11 +25,30 @@ from keras_nlp.utils.python_utils import format_docstring
 class Backbone(keras.Model):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self._token_embedding = None
+
+    def __setattr__(self, name, value):
+        # Work around torch setattr for properties.
+        if name in ["token_embedding"]:
+            return object.__setattr__(self, name, value)
+        return super().__setattr__(name, value)
 
     @property
     def token_embedding(self):
-        """A `keras.layers.Embedding` instance for embedding token ids."""
-        raise NotImplementedError
+        """A `keras.layers.Embedding` instance for embedding token ids.
+
+        This layer integer token ids to the hidden dim of the model.
+        """
+        return self._token_embedding
+
+    @token_embedding.setter
+    def token_embedding(self, value):
+        # Workaround tf.keras h5 checkpoint loading, which is sensitive to layer
+        # count mismatches and does not deduplicate layers. This could go away
+        # if we update our checkpoints to the newer `.weights.h5` format.
+        self._setattr_tracking = False
+        self._token_embedding = value
+        self._setattr_tracking = True
 
     def get_config(self):
         # Don't chain to super here. The default `get_config()` for functional
diff --git a/keras_nlp/models/bart/bart_backbone.py b/keras_nlp/models/bart/bart_backbone.py
index 594f2467..ae26a919 100644
--- a/keras_nlp/models/bart/bart_backbone.py
+++ b/keras_nlp/models/bart/bart_backbone.py
@@ -19,6 +19,7 @@ import copy
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.layers.modeling.transformer_decoder import TransformerDecoder
 from keras_nlp.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_nlp.models.backbone import Backbone
@@ -120,7 +121,7 @@ class BartBackbone(Backbone):
         )
 
         # Token embedding layer. This layer is shared by encoder and decoder.
-        token_embedding_layer = keras.layers.Embedding(
+        token_embedding_layer = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=bart_kernel_initializer(),
@@ -240,6 +241,7 @@ class BartBackbone(Backbone):
         self.intermediate_dim = intermediate_dim
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
+        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
@@ -257,10 +259,6 @@ class BartBackbone(Backbone):
 
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_embedding")
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/bart/bart_seq_2_seq_lm.py b/keras_nlp/models/bart/bart_seq_2_seq_lm.py
index 1fde9309..a4f98d21 100644
--- a/keras_nlp/models/bart/bart_seq_2_seq_lm.py
+++ b/keras_nlp/models/bart/bart_seq_2_seq_lm.py
@@ -27,20 +27,6 @@ from keras_nlp.models.generative_task import GenerativeTask
 from keras_nlp.utils.python_utils import classproperty
 
 
-# TODO: Extend and factor this out into keras_nlp.layers.
-class ReverseEmbedding(keras.layers.Layer):
-    def __init__(self, embedding, **kwargs):
-        super().__init__(**kwargs)
-        self.embedding = embedding
-
-    def call(self, inputs):
-        kernel = ops.transpose(ops.convert_to_tensor(self.embedding.embeddings))
-        return ops.matmul(inputs, kernel)
-
-    def compute_output_shape(self, input_shape):
-        return (input_shape[0],) + (self.embedding.embeddings.shape[0],)
-
-
 @keras_nlp_export("keras_nlp.models.BartSeq2SeqLM")
 class BartSeq2SeqLM(GenerativeTask):
     """An end-to-end BART model for seq2seq language modeling.
@@ -201,13 +187,8 @@ class BartSeq2SeqLM(GenerativeTask):
         **kwargs,
     ):
         inputs = backbone.input
-        x = backbone(inputs)["decoder_sequence_output"]
-        # Use token embedding weights to project from the token representation
-        # to vocabulary logits.
-        outputs = ReverseEmbedding(
-            backbone.token_embedding,
-            name="reverse_embedding",
-        )(x)
+        hidden_states = backbone(inputs)["decoder_sequence_output"]
+        outputs = backbone.token_embedding(hidden_states, reverse=True)
 
         # Instantiate using Functional API Model constructor.
         super().__init__(
@@ -347,7 +328,7 @@ class BartSeq2SeqLM(GenerativeTask):
             cross_attention_cache = ops.stack(cross_attention_caches, axis=1)
 
         hidden_states = x
-        logits = self.get_layer("reverse_embedding")(x)
+        logits = self.backbone.token_embedding(hidden_states, reverse=True)
         return (
             logits,
             hidden_states,
diff --git a/keras_nlp/models/bert/bert_backbone.py b/keras_nlp/models/bert/bert_backbone.py
index 8bb6057a..155da972 100644
--- a/keras_nlp/models/bert/bert_backbone.py
+++ b/keras_nlp/models/bert/bert_backbone.py
@@ -19,6 +19,7 @@ import copy
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.bert.bert_presets import backbone_presets
@@ -113,7 +114,7 @@ class BertBackbone(Backbone):
         )
 
         # Embed tokens, positions, and segment ids.
-        token_embedding_layer = keras.layers.Embedding(
+        token_embedding_layer = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=bert_kernel_initializer(),
@@ -196,6 +197,7 @@ class BertBackbone(Backbone):
         self.max_sequence_length = max_sequence_length
         self.num_segments = num_segments
         self.cls_token_index = cls_token_index
+        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
@@ -213,10 +215,6 @@ class BertBackbone(Backbone):
         )
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_embedding")
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/bert/bert_masked_lm.py b/keras_nlp/models/bert/bert_masked_lm.py
index f99314e9..d696f77f 100644
--- a/keras_nlp/models/bert/bert_masked_lm.py
+++ b/keras_nlp/models/bert/bert_masked_lm.py
@@ -111,7 +111,7 @@ class BertMaskedLM(Task):
         backbone_outputs = backbone(backbone.input)
         outputs = MaskedLMHead(
             vocabulary_size=backbone.vocabulary_size,
-            embedding_weights=backbone.token_embedding.embeddings,
+            token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=bert_kernel_initializer(),
             name="mlm_head",
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
index c4bfe946..43a5a34c 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
@@ -18,6 +18,7 @@ import copy
 
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.deberta_v3.deberta_v3_presets import backbone_presets
 from keras_nlp.models.deberta_v3.disentangled_attention_encoder import (
@@ -118,12 +119,13 @@ class DebertaV3Backbone(Backbone):
         )
 
         # Embed tokens.
-        x = keras.layers.Embedding(
+        token_embedding_layer = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=deberta_kernel_initializer(),
             name="token_embedding",
-        )(token_id_input)
+        )
+        x = token_embedding_layer(token_id_input)
 
         # Normalize and apply dropout to embeddings.
         x = keras.layers.LayerNormalization(
@@ -184,6 +186,7 @@ class DebertaV3Backbone(Backbone):
         self.max_sequence_length = max_sequence_length
         self.bucket_size = bucket_size
         self.start_token_index = 0
+        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
@@ -201,10 +204,6 @@ class DebertaV3Backbone(Backbone):
         )
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_embedding")
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py
index d79453a4..2fcb8bf4 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py
@@ -114,7 +114,7 @@ class DebertaV3MaskedLM(Task):
         backbone_outputs = backbone(backbone.input)
         outputs = MaskedLMHead(
             vocabulary_size=backbone.vocabulary_size,
-            embedding_weights=backbone.token_embedding.embeddings,
+            token_embedding=backbone.token_embedding,
             intermediate_activation=lambda x: keras.activations.gelu(
                 x, approximate=False
             ),
diff --git a/keras_nlp/models/distil_bert/distil_bert_backbone.py b/keras_nlp/models/distil_bert/distil_bert_backbone.py
index 23f9cfbf..a00b15da 100644
--- a/keras_nlp/models/distil_bert/distil_bert_backbone.py
+++ b/keras_nlp/models/distil_bert/distil_bert_backbone.py
@@ -111,13 +111,14 @@ class DistilBertBackbone(Backbone):
         )
 
         # Embed tokens and positions.
-        x = TokenAndPositionEmbedding(
+        embedding_layer = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
             sequence_length=max_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=distilbert_kernel_initializer(),
             name="token_and_position_embedding",
-        )(token_id_input)
+        )
+        x = embedding_layer(token_id_input)
 
         # Normalize and apply dropout to embeddings.
         x = keras.layers.LayerNormalization(
@@ -161,6 +162,7 @@ class DistilBertBackbone(Backbone):
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
         self.cls_token_index = 0
+        self.token_embedding = embedding_layer.token_embedding
 
     def get_config(self):
         config = super().get_config()
@@ -177,10 +179,6 @@ class DistilBertBackbone(Backbone):
         )
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_and_position_embedding").token_embedding
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/distil_bert/distil_bert_masked_lm.py b/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
index b307ea10..0f5299a5 100644
--- a/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
+++ b/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
@@ -114,7 +114,7 @@ class DistilBertMaskedLM(Task):
         backbone_outputs = backbone(backbone.input)
         outputs = MaskedLMHead(
             vocabulary_size=backbone.vocabulary_size,
-            embedding_weights=backbone.token_embedding.embeddings,
+            token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=distilbert_kernel_initializer(),
             name="mlm_head",
diff --git a/keras_nlp/models/f_net/f_net_backbone.py b/keras_nlp/models/f_net/f_net_backbone.py
index 166dae56..5053a55d 100644
--- a/keras_nlp/models/f_net/f_net_backbone.py
+++ b/keras_nlp/models/f_net/f_net_backbone.py
@@ -20,6 +20,7 @@ from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.f_net_encoder import FNetEncoder
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.f_net.f_net_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
@@ -112,7 +113,7 @@ class FNetBackbone(Backbone):
         )
 
         # Embed tokens, positions, and segment ids.
-        token_embedding_layer = keras.layers.Embedding(
+        token_embedding_layer = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=f_net_kernel_initializer(),
@@ -200,6 +201,7 @@ class FNetBackbone(Backbone):
         self.max_sequence_length = max_sequence_length
         self.num_segments = num_segments
         self.cls_token_index = cls_token_index
+        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
@@ -216,10 +218,6 @@ class FNetBackbone(Backbone):
         )
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_embedding")
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/f_net/f_net_masked_lm.py b/keras_nlp/models/f_net/f_net_masked_lm.py
index cc8d44eb..0355a235 100644
--- a/keras_nlp/models/f_net/f_net_masked_lm.py
+++ b/keras_nlp/models/f_net/f_net_masked_lm.py
@@ -109,7 +109,7 @@ class FNetMaskedLM(Task):
         backbone_outputs = backbone(backbone.input)
         outputs = MaskedLMHead(
             vocabulary_size=backbone.vocabulary_size,
-            embedding_weights=backbone.token_embedding.embeddings,
+            token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=f_net_kernel_initializer(),
             name="mlm_head",
diff --git a/keras_nlp/models/gpt2/gpt2_backbone.py b/keras_nlp/models/gpt2/gpt2_backbone.py
index 6a00e6ec..faf12c73 100644
--- a/keras_nlp/models/gpt2/gpt2_backbone.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone.py
@@ -23,6 +23,7 @@ from tensorflow.keras.dtensor.experimental import LayoutMap
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
 from keras_nlp.layers.modeling.position_embedding import PositionEmbedding
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.layers.modeling.transformer_decoder import TransformerDecoder
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.gpt2.gpt2_presets import backbone_presets
@@ -108,12 +109,13 @@ class GPT2Backbone(Backbone):
         )
 
         # Embed tokens, positions.
-        token_embedding = keras.layers.Embedding(
+        token_embedding_layer = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=_gpt_2_kernel_initializer(stddev=0.01),
             name="token_embedding",
-        )(token_ids)
+        )
+        token_embedding = token_embedding_layer(token_ids)
 
         # Can't use `TokenAndPositionEmbedding` layer here because of different
         # initializers.
@@ -171,6 +173,7 @@ class GPT2Backbone(Backbone):
         self.intermediate_dim = intermediate_dim
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
+        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
@@ -187,10 +190,6 @@ class GPT2Backbone(Backbone):
         )
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_embedding")
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm.py b/keras_nlp/models/gpt2/gpt2_causal_lm.py
index 19de4a4e..6e508c43 100644
--- a/keras_nlp/models/gpt2/gpt2_causal_lm.py
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm.py
@@ -27,20 +27,6 @@ from keras_nlp.models.gpt2.gpt2_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
 
 
-# TODO: Extend and factor this out into keras_nlp.layers.
-class ReverseEmbedding(keras.layers.Layer):
-    def __init__(self, embedding, **kwargs):
-        super().__init__(**kwargs)
-        self.embedding = embedding
-
-    def call(self, inputs):
-        kernel = ops.transpose(ops.convert_to_tensor(self.embedding.embeddings))
-        return ops.matmul(inputs, kernel)
-
-    def compute_output_shape(self, input_shape):
-        return (input_shape[0],) + (self.embedding.embeddings.shape[0],)
-
-
 @keras_nlp_export("keras_nlp.models.GPT2CausalLM")
 class GPT2CausalLM(GenerativeTask):
     """An end-to-end GPT2 model for causal langauge modeling.
@@ -171,13 +157,8 @@ class GPT2CausalLM(GenerativeTask):
         **kwargs,
     ):
         inputs = backbone.input
-        x = backbone(inputs)
-        # Use token embedding weights to project from the token representation
-        # to vocabulary logits.
-        outputs = ReverseEmbedding(
-            backbone.token_embedding,
-            name="reverse_embedding",
-        )(x)
+        hidden_states = backbone(inputs)
+        outputs = backbone.token_embedding(hidden_states, reverse=True)
 
         # Instantiate using Functional API Model constructor.
         super().__init__(
@@ -255,9 +236,8 @@ class GPT2CausalLM(GenerativeTask):
             )
             caches.append(next_cache)
         cache = ops.stack(caches, axis=1)
-        x = self.backbone.get_layer("layer_norm")(x)
-        hidden_states = x
-        logits = self.get_layer("reverse_embedding")(x)
+        hidden_states = x = self.backbone.get_layer("layer_norm")(x)
+        logits = self.backbone.get_layer("token_embedding")(x, reverse=True)
         return logits, hidden_states, cache
 
     def _build_cache(self, token_ids):
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
index 43689c27..2919976d 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
@@ -16,6 +16,7 @@
 
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.gpt_neo_x.gpt_neo_x_decoder import GPTNeoXDecoder
 
@@ -84,12 +85,13 @@ class GPTNeoXBackbone(Backbone):
         )
 
         # Embed tokens
-        token_embedding = keras.layers.Embedding(
+        token_embedding_layer = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=_gpt_neo_x_kernel_initializer(stddev=0.01),
             name="token_embedding",
-        )(token_ids)
+        )
+        token_embedding = token_embedding_layer(token_ids)
 
         x = keras.layers.Dropout(
             dropout,
@@ -140,6 +142,7 @@ class GPTNeoXBackbone(Backbone):
         self.rotary_max_wavelength = rotary_max_wavelength
         self.max_sequence_length = max_sequence_length
         self.layer_norm_epsilon = layer_norm_epsilon
+        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
@@ -158,7 +161,3 @@ class GPTNeoXBackbone(Backbone):
             }
         )
         return config
-
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_embedding")
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm.py
index f74aa1ae..6462a6a8 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm.py
@@ -24,19 +24,6 @@ from keras_nlp.models.gpt_neo_x.gpt_neo_x_causal_lm_preprocessor import (
 from keras_nlp.utils.python_utils import classproperty
 
 
-class ReverseEmbedding(keras.layers.Layer):
-    def __init__(self, embedding, **kwargs):
-        super().__init__(**kwargs)
-        self.embedding = embedding
-
-    def call(self, inputs):
-        kernel = ops.transpose(ops.convert_to_tensor(self.embedding.embeddings))
-        return ops.matmul(inputs, kernel)
-
-    def compute_output_shape(self, input_shape):
-        return (input_shape[0],) + (self.embedding.embeddings.shape[0],)
-
-
 @keras_nlp_export("keras_nlp.models.GPTNeoXCausalLM")
 class GPTNeoXCausalLM(GenerativeTask):
     """An end-to-end GPTNeoX model for causal langauge modeling.
@@ -67,13 +54,8 @@ class GPTNeoXCausalLM(GenerativeTask):
         **kwargs,
     ):
         inputs = backbone.input
-        x = backbone(inputs)
-        # Use token embedding weights to project from the token representation
-        # to vocabulary logits.
-        outputs = ReverseEmbedding(
-            backbone.token_embedding,
-            name="reverse_embedding",
-        )(x)
+        hidden_states = backbone(inputs)
+        outputs = backbone.token_embedding(hidden_states, reverse=True)
 
         # Instantiate using Functional API Model constructor.
         super().__init__(
@@ -143,7 +125,7 @@ class GPTNeoXCausalLM(GenerativeTask):
         cache = ops.stack(caches, axis=1)
         x = self.backbone.get_layer("layer_norm")(x)
         hidden_states = x
-        logits = self.get_layer("reverse_embedding")(x)
+        logits = self.backbone.token_embedding(hidden_states, reverse=True)
         return logits, hidden_states, cache
 
     def _build_cache(self, token_ids):
diff --git a/keras_nlp/models/opt/opt_backbone.py b/keras_nlp/models/opt/opt_backbone.py
index 996511f6..4052644c 100644
--- a/keras_nlp/models/opt/opt_backbone.py
+++ b/keras_nlp/models/opt/opt_backbone.py
@@ -106,13 +106,14 @@ class OPTBackbone(Backbone):
         )
 
         # Embed tokens and positions.
-        x = TokenAndPositionEmbedding(
+        embedding_layer = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
             sequence_length=max_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=opt_kernel_initializer(),
             name="embeddings",
-        )(token_ids)
+        )
+        x = embedding_layer(token_ids)
 
         # Apply successive transformer decoder blocks.
         for i in range(num_layers):
@@ -153,6 +154,7 @@ class OPTBackbone(Backbone):
         self.intermediate_dim = intermediate_dim
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
+        self.token_embedding = embedding_layer.token_embedding
 
     def get_config(self):
         return {
@@ -165,10 +167,6 @@ class OPTBackbone(Backbone):
             "max_sequence_length": self.max_sequence_length,
         }
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("embeddings").token_embedding
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/opt/opt_causal_lm.py b/keras_nlp/models/opt/opt_causal_lm.py
index fccaf4ce..c649da85 100644
--- a/keras_nlp/models/opt/opt_causal_lm.py
+++ b/keras_nlp/models/opt/opt_causal_lm.py
@@ -27,20 +27,6 @@ from keras_nlp.models.opt.opt_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
 
 
-# TODO: Extend and factor this out into keras_nlp.layers.
-class ReverseEmbedding(keras.layers.Layer):
-    def __init__(self, embedding, **kwargs):
-        super().__init__(**kwargs)
-        self.embedding = embedding
-
-    def call(self, inputs):
-        kernel = ops.transpose(ops.convert_to_tensor(self.embedding.embeddings))
-        return ops.matmul(inputs, kernel)
-
-    def compute_output_shape(self, input_shape):
-        return (input_shape[0],) + (self.embedding.embeddings.shape[0],)
-
-
 @keras_nlp_export("keras_nlp.models.OPTCausalLM")
 class OPTCausalLM(GenerativeTask):
     """An end-to-end OPT model for causal langauge modeling.
@@ -171,13 +157,8 @@ class OPTCausalLM(GenerativeTask):
         **kwargs,
     ):
         inputs = backbone.input
-        x = backbone(inputs)
-        # Use token embedding weights to project from the token representation
-        # to vocabulary logits.
-        outputs = ReverseEmbedding(
-            backbone.token_embedding,
-            name="reverse_embedding",
-        )(x)
+        hidden_states = backbone(inputs)
+        outputs = backbone.token_embedding(hidden_states, reverse=True)
 
         # Instantiate using Functional API Model constructor.
         super().__init__(
@@ -252,7 +233,7 @@ class OPTCausalLM(GenerativeTask):
         cache = ops.stack(caches, axis=1)
         x = self.backbone.get_layer("layer_norm")(x)
         hidden_states = x
-        logits = self.get_layer("reverse_embedding")(x)
+        logits = self.backbone.token_embedding(hidden_states, reverse=True)
         return logits, hidden_states, cache
 
     def _build_cache(self, token_ids):
diff --git a/keras_nlp/models/preprocessor.py b/keras_nlp/models/preprocessor.py
index f24b9203..fc2604b3 100644
--- a/keras_nlp/models/preprocessor.py
+++ b/keras_nlp/models/preprocessor.py
@@ -31,9 +31,8 @@ class Preprocessor(PreprocessingLayer):
     def __setattr__(self, name, value):
         # Work around torch setattr for properties.
         if name in ["tokenizer"]:
-            object.__setattr__(self, name, value)
-        else:
-            super().__setattr__(name, value)
+            return object.__setattr__(self, name, value)
+        return super().__setattr__(name, value)
 
     @property
     def tokenizer(self):
diff --git a/keras_nlp/models/roberta/roberta_backbone.py b/keras_nlp/models/roberta/roberta_backbone.py
index f8440b0c..29bc1405 100644
--- a/keras_nlp/models/roberta/roberta_backbone.py
+++ b/keras_nlp/models/roberta/roberta_backbone.py
@@ -160,6 +160,7 @@ class RobertaBackbone(Backbone):
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
         self.start_token_index = 0
+        self.token_embedding = embedding_layer.token_embedding
 
     def get_config(self):
         config = super().get_config()
@@ -176,10 +177,6 @@ class RobertaBackbone(Backbone):
         )
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("embeddings").token_embedding
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/roberta/roberta_masked_lm.py b/keras_nlp/models/roberta/roberta_masked_lm.py
index 37f5dc39..f80566c3 100644
--- a/keras_nlp/models/roberta/roberta_masked_lm.py
+++ b/keras_nlp/models/roberta/roberta_masked_lm.py
@@ -113,7 +113,7 @@ class RobertaMaskedLM(Task):
         backbone_outputs = backbone(backbone.input)
         outputs = MaskedLMHead(
             vocabulary_size=backbone.vocabulary_size,
-            embedding_weights=backbone.token_embedding.embeddings,
+            token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=roberta_kernel_initializer(),
             name="mlm_head",
diff --git a/keras_nlp/models/t5/t5_backbone.py b/keras_nlp/models/t5/t5_backbone.py
index b23e0bec..071405f7 100644
--- a/keras_nlp/models/t5/t5_backbone.py
+++ b/keras_nlp/models/t5/t5_backbone.py
@@ -16,6 +16,7 @@
 
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.backend import keras
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.t5.t5_layer_norm import T5LayerNorm
 from keras_nlp.models.t5.t5_transformer_layer import T5TransformerLayer
@@ -63,6 +64,9 @@ class T5Backbone(Backbone):
             recent versions do. Defaults to `True`.
         layer_norm_epsilon: float. Epsilon factor to be used in the
             layer normalization layers in the Transformer layers.
+        tie_embedding_weights: boolean. If `True`, the weights of the token
+            embedding and the weights projecting language model outputs from
+            `hidden_dim`
     """
 
     def __init__(
@@ -76,6 +80,7 @@ class T5Backbone(Backbone):
         activation="gelu",
         use_gated_activation=True,
         layer_norm_epsilon=1e-06,
+        tie_embedding_weights=False,
         **kwargs,
     ):
         assert_tf_backend(self.__class__.__name__)
@@ -97,9 +102,10 @@ class T5Backbone(Backbone):
         )
 
         # Token embedding layer. This layer is shared by encoder and decoder.
-        token_embedding_layer = keras.layers.Embedding(
+        token_embedding_layer = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
+            tie_weights=tie_embedding_weights,
             embeddings_initializer=keras.initializers.TruncatedNormal(1.0),
             name="token_embedding",
         )
@@ -210,6 +216,8 @@ class T5Backbone(Backbone):
         self.activation = keras.activations.get(activation)
         self.dropout = dropout
         self.layer_norm_epsilon = layer_norm_epsilon
+        self.tie_embedding_weights = tie_embedding_weights
+        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
@@ -223,14 +231,11 @@ class T5Backbone(Backbone):
                 "activation": keras.activations.serialize(self.activation),
                 "dropout": self.dropout,
                 "layer_norm_epsilon": self.layer_norm_epsilon,
+                "tie_embedding_weights": self.tie_embedding_weights,
             }
         )
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_embedding")
-
     @classproperty
     def presets(cls):
         return {}
diff --git a/keras_nlp/models/task.py b/keras_nlp/models/task.py
index e8dd309e..583c1072 100644
--- a/keras_nlp/models/task.py
+++ b/keras_nlp/models/task.py
@@ -89,9 +89,8 @@ class Task(PipelineModel):
     def __setattr__(self, name, value):
         # Work around torch setattr for properties.
         if name in ["backbone", "preprocessor"]:
-            object.__setattr__(self, name, value)
-        else:
-            super().__setattr__(name, value)
+            return object.__setattr__(self, name, value)
+        return super().__setattr__(name, value)
 
     @property
     def backbone(self):
diff --git a/keras_nlp/models/whisper/whisper_backbone.py b/keras_nlp/models/whisper/whisper_backbone.py
index 1bd104d5..5bdb219e 100644
--- a/keras_nlp/models/whisper/whisper_backbone.py
+++ b/keras_nlp/models/whisper/whisper_backbone.py
@@ -209,13 +209,14 @@ class WhisperBackbone(Backbone):
         # ====== Decoder ======
 
         # Embed tokens and positions.
-        x = TokenAndPositionEmbedding(
+        embedding_layer = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
             sequence_length=max_decoder_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=whisper_kernel_initializer(),
             name="decoder_token_and_position_embedding",
-        )(decoder_token_id_input)
+        )
+        x = embedding_layer(decoder_token_id_input)
 
         # Apply dropout to embeddings.
         x = keras.layers.Dropout(
@@ -275,6 +276,7 @@ class WhisperBackbone(Backbone):
         self.dropout = dropout
         self.max_encoder_sequence_length = max_encoder_sequence_length
         self.max_decoder_sequence_length = max_decoder_sequence_length
+        self.token_embedding = embedding_layer
 
     def get_config(self):
         config = super().get_config()
@@ -293,12 +295,6 @@ class WhisperBackbone(Backbone):
         )
         return config
 
-    @property
-    def token_embedding(self):
-        return self.get_layer(
-            "decoder_token_and_position_embedding"
-        ).token_embedding
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py b/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py
index 80715624..b0d9b7fc 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py
@@ -116,7 +116,7 @@ class XLMRobertaMaskedLM(Task):
         backbone_outputs = backbone(backbone.input)
         outputs = MaskedLMHead(
             vocabulary_size=backbone.vocabulary_size,
-            embedding_weights=backbone.token_embedding.embeddings,
+            token_embedding=backbone.token_embedding,
             intermediate_activation="gelu",
             kernel_initializer=roberta_kernel_initializer(),
             name="mlm_head",


Test patch:
diff --git a/keras_nlp/layers/modeling/masked_lm_head_test.py b/keras_nlp/layers/modeling/masked_lm_head_test.py
index 7d85da5b..62532967 100644
--- a/keras_nlp/layers/modeling/masked_lm_head_test.py
+++ b/keras_nlp/layers/modeling/masked_lm_head_test.py
@@ -17,13 +17,14 @@ import os
 
 from keras_nlp.backend import keras
 from keras_nlp.backend import ops
-from keras_nlp.layers.modeling import masked_lm_head
+from keras_nlp.layers.modeling.masked_lm_head import MaskedLMHead
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
 from keras_nlp.tests.test_case import TestCase
 
 
 class MaskedLMHeadTest(TestCase):
     def test_valid_call(self):
-        head = masked_lm_head.MaskedLMHead(
+        head = MaskedLMHead(
             vocabulary_size=100,
             activation="softmax",
         )
@@ -36,12 +37,12 @@ class MaskedLMHeadTest(TestCase):
         position_data = ops.random.randint(minval=0, maxval=10, shape=(4, 5))
         model((token_data, position_data))
 
-    def test_valid_call_with_embedding_weights(self):
-        embedding = keras.layers.Embedding(100, 16)
+    def test_valid_call_with_token_embedding(self):
+        embedding = ReversibleEmbedding(100, 16)
         embedding.build((4, 10))
-        head = masked_lm_head.MaskedLMHead(
+        head = MaskedLMHead(
             vocabulary_size=100,
-            embedding_weights=embedding.embeddings,
+            token_embedding=embedding,
             activation="softmax",
         )
         # Use a difference "hidden dim" for the model than "embedding dim", we
@@ -55,7 +56,7 @@ class MaskedLMHeadTest(TestCase):
         model((sequence_data, position_data))
 
     def test_get_config_and_from_config(self):
-        head = masked_lm_head.MaskedLMHead(
+        head = MaskedLMHead(
             vocabulary_size=100,
             kernel_initializer="HeNormal",
             bias_initializer="Zeros",
@@ -79,7 +80,7 @@ class MaskedLMHeadTest(TestCase):
 
         self.assertEqual(config, {**config, **expected_params})
 
-        restored = masked_lm_head.MaskedLMHead.from_config(config)
+        restored = MaskedLMHead.from_config(config)
         restored_config = restored.get_config()
 
         self.assertEqual(
@@ -89,19 +90,19 @@ class MaskedLMHeadTest(TestCase):
 
     def test_value_error_when_neither_embedding_or_vocab_size_set(self):
         with self.assertRaises(ValueError):
-            masked_lm_head.MaskedLMHead()
+            MaskedLMHead()
 
     def test_value_error_when_vocab_size_mismatch(self):
-        embedding = keras.layers.Embedding(100, 16)
+        embedding = ReversibleEmbedding(100, 16)
         embedding.build((4, 10))
         with self.assertRaises(ValueError):
-            masked_lm_head.MaskedLMHead(
+            MaskedLMHead(
                 vocabulary_size=101,
-                embedding_weights=embedding.embeddings,
+                token_embedding=embedding,
             )
 
     def test_one_train_step(self):
-        head = masked_lm_head.MaskedLMHead(
+        head = MaskedLMHead(
             vocabulary_size=100,
         )
         encoded_tokens = keras.Input(shape=(10, 16))
@@ -120,7 +121,7 @@ class MaskedLMHeadTest(TestCase):
         self.assertGreater(loss, 0)
 
     def test_saved_model(self):
-        head = masked_lm_head.MaskedLMHead(
+        head = MaskedLMHead(
             vocabulary_size=100,
             activation="softmax",
         )
diff --git a/keras_nlp/layers/modeling/reversible_embedding_test.py b/keras_nlp/layers/modeling/reversible_embedding_test.py
new file mode 100644
index 00000000..e1de16d1
--- /dev/null
+++ b/keras_nlp/layers/modeling/reversible_embedding_test.py
@@ -0,0 +1,101 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import numpy as np
+from absl.testing import parameterized
+
+from keras_nlp.backend import keras
+from keras_nlp.backend import ops
+from keras_nlp.layers.modeling.reversible_embedding import ReversibleEmbedding
+from keras_nlp.tests.test_case import TestCase
+
+
+class ReversibleEmbeddingTest(TestCase):
+    @parameterized.named_parameters(
+        ("tie_weights", True),
+        ("untie_weights", False),
+    )
+    def test_valid_call(self, tie_weights):
+        embedding = ReversibleEmbedding(100, 32, tie_weights=tie_weights)
+        inputs = keras.Input(shape=(10,))
+        hidden_states = embedding(inputs)
+        outputs = embedding(hidden_states, reverse=True)
+        model = keras.Model(inputs, outputs)
+
+        input_data = ops.random.uniform(shape=(4, 10))
+        model(input_data)
+
+    def test_correctness(self):
+        layer = ReversibleEmbedding(input_dim=3, output_dim=2)
+        layer.build()
+        layer.embeddings.assign(np.array([[0.0, 0.0], [2.0, 2.0], [3.0, 3.0]]))
+        out = layer(np.array(([2, 1, 0])))
+        self.assertAllClose(out, np.array([[3.0, 3.0], [2.0, 2.0], [0.0, 0.0]]))
+
+        layer = ReversibleEmbedding(input_dim=3, output_dim=2)
+        layer.build()
+        layer.embeddings.assign(np.array([[0.0, 0.0], [2.0, 2.0], [3.0, 3.0]]))
+        out = layer(np.array(([[1.0, 1.0]])), reverse=True)
+        self.assertAllClose(out, np.array([[0.0, 4.0, 6.0]]))
+
+    def test_config(self):
+        original = ReversibleEmbedding(
+            100,
+            32,
+            tie_weights=False,
+            embeddings_initializer="HeNormal",
+        )
+        restored = ReversibleEmbedding.from_config(original.get_config())
+        restored.set_weights(original.get_weights())
+        self.assertEqual(restored.get_config(), original.get_config())
+
+    def test_tied_checkpoint_untied_weights(self):
+        embedding = ReversibleEmbedding(100, 16, tie_weights=True)
+        inputs = keras.Input(shape=(10,), dtype="int32")
+        hidden_states = embedding(inputs)
+        outputs = embedding(hidden_states, reverse=True)
+        tied_model = keras.Model(inputs, outputs)
+        path = os.path.join(self.get_temp_dir(), "checkpoint.weights.h5")
+        tied_model.save_weights(path)
+
+        embedding = ReversibleEmbedding(100, 16, tie_weights=False)
+        inputs = keras.Input(shape=(10,), dtype="int32")
+        hidden_states = embedding(inputs)
+        outputs = embedding(hidden_states, reverse=True)
+        untied_model = keras.Model(inputs, outputs)
+        untied_model.load_weights(path)
+
+        input_data = ops.ones(shape=(4, 10), dtype="int32")
+        self.assertAllClose(untied_model(input_data), tied_model(input_data))
+
+    @parameterized.named_parameters(
+        ("tie_weights", True),
+        ("untie_weights", False),
+    )
+    def test_saved_model(self, tie_weights):
+        embedding = ReversibleEmbedding(100, 16, tie_weights=tie_weights)
+        inputs = keras.Input(shape=(10,), dtype="int32")
+        hidden_states = embedding(inputs)
+        outputs = embedding(hidden_states, reverse=True)
+        model = keras.Model(inputs, outputs)
+
+        input_data = ops.ones(shape=(4, 10), dtype="int32")
+        model_output = model(input_data)
+        path = os.path.join(self.get_temp_dir(), "model.keras")
+        model.save(path, save_format="keras_v3")
+        restored_model = keras.models.load_model(path)
+
+        restored_output = restored_model(input_data)
+        self.assertAllClose(model_output, restored_output)