Add soft capping to reversible embedding layer (#1718)

Forgetting the final output soft-cap is a really easy mistake,
and worse, outputs will still look plausible for generations without
the softcap, just with worse actual results.

Adding this to our reversible embedding layer will be much more robust.
As long as you use the layer to compute logits over the vocab, you can
no longer forget the soft-cap.

Author: mattdangerw

keras_nlp/src/layers/modeling/reversible_embedding.py

@@ -50,6 +50,10 @@ class ReversibleEmbedding(keras.layers.Embedding):
             "padding" value that should be masked out.
         reverse_dtype: The dtype for the reverse projection computation.
             Defaults to the `compute_dtype` of the layer.
+        logit_soft_cap: If `logit_soft_cap` is set and `reverse=True`, the
+            output logits will be scaled by
+            `tanh(logits / logit_soft_cap) * logit_soft_cap`. This narrows the
+            range of output logits and can improve training.
         **kwargs: other keyword arguments passed to `keras.layers.Embedding`,
             including `name`, `trainable`, `dtype` etc.
 
@@ -91,6 +95,7 @@ def __init__(
         embeddings_constraint=None,
         mask_zero=False,
         reverse_dtype=None,
+        logit_soft_cap=None,
         **kwargs,
     ):
         super().__init__(
@@ -104,6 +109,7 @@ def __init__(
         )
         self.tie_weights = tie_weights
         self.reverse_dtype = reverse_dtype
+        self.logit_soft_cap = logit_soft_cap
 
     def build(self, inputs_shape=None):
         super().build(inputs_shape)
@@ -125,7 +131,12 @@ def call(self, inputs, reverse=False):
             if self.reverse_dtype is not None:
                 inputs = ops.cast(inputs, self.reverse_dtype)
                 kernel = ops.cast(kernel, self.reverse_dtype)
-            return ops.matmul(inputs, kernel)
+            logits = ops.matmul(inputs, kernel)
+            # Optionally soft-cap logits.
+            if self.logit_soft_cap is not None:
+                soft_cap = self.logit_soft_cap
+                logits = ops.tanh(logits / soft_cap) * soft_cap
+            return logits
 
         return super().call(inputs)
 
@@ -135,6 +146,7 @@ def get_config(self):
             {
                 "tie_weights": self.tie_weights,
                 "reverse_dtype": self.reverse_dtype,
+                "logit_soft_cap": self.logit_soft_cap,
             }
         )
         return config

keras_nlp/src/layers/modeling/reversible_embedding_test.py

@@ -39,6 +39,7 @@ def test_layer_behaviors_tied(self, tie_weights):
                 "output_dim": 32,
                 "tie_weights": tie_weights,
                 "embeddings_initializer": "HeNormal",
+                "logit_soft_cap": 50,
             },
             input_data=random.randint(minval=0, maxval=100, shape=(4, 10)),
             expected_output_shape=(4, 10, 32),
@@ -80,6 +81,12 @@ def test_correctness(self):
         out = layer(np.array(([[1.0, 1.0]])), reverse=True)
         self.assertAllClose(out, np.array([[0.0, 4.0, 6.0]]))
 
+        layer = ReversibleEmbedding(input_dim=3, output_dim=2, logit_soft_cap=5)
+        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, 3.320184, 4.168273]]))
+
     def test_reverse_dtype(self):
         embedding = ReversibleEmbedding(100, 16, reverse_dtype="float32")
         input_data = ops.ones(shape=(4, 10, 16))

keras_nlp/src/models/gemma/gemma_backbone.py

@@ -132,6 +132,7 @@ def __init__(
                 seed=None,
             ),
             dtype=dtype,
+            logit_soft_cap=final_logit_soft_cap,
             name="token_embedding",
         )
         self.transformer_layers = []

keras_nlp/src/models/gemma/gemma_causal_lm.py

@@ -227,13 +227,6 @@ def call_with_cache(
         cache = ops.stack(caches, axis=1)
         hidden_states = x = self.backbone.layer_norm(x)
         logits = self.backbone.token_embedding(x, reverse=True)
-
-        if self.backbone.final_logit_soft_cap is not None:
-            logits = ops.divide(logits, self.backbone.final_logit_soft_cap)
-            logits = ops.multiply(
-                ops.tanh(logits), self.backbone.final_logit_soft_cap
-            )
-
         return logits, hidden_states, cache
 
     def _build_cache(self, token_ids):
@@ -445,12 +438,6 @@ def default_layer_intercept_fn(x, unused_i):
         x = self.backbone.layer_norm(x)
         logits = self.backbone.token_embedding(x, reverse=True)
 
-        if self.backbone.final_logit_soft_cap is not None:
-            logits = ops.divide(logits, self.backbone.final_logit_soft_cap)
-            logits = ops.multiply(
-                ops.tanh(logits), self.backbone.final_logit_soft_cap
-            )
-
         if scoring_mode == "logits":
             return logits