Fix attention score returns in CachedMultiHeadAttention (#2799)

* Fix attention score returns in cached MHA

* Address cached attention test feedback

Author: GargiGupta-io

keras_hub/src/layers/modeling/cached_multi_head_attention.py

@@ -52,16 +52,17 @@ class CachedMultiHeadAttention(keras.layers.MultiHeadAttention):
             `cache` (usually the index of the current token being processed
             when running generation). If `cache_update_index=None` while `cache`
             is set, the cache will not be updated.
+        return_attention_scores: a boolean indicating whether the output
+            should include attention scores.
         training: a boolean indicating whether the layer should behave in
             training mode or in inference mode.
 
     Returns:
-        An `(attention_output, cache)` tuple. `attention_output` is the result
-        of the computation, of shape `(B, T, dim)`, where `T` is for target
-        sequence shapes and `dim` is the query input last dimension if
-        `output_shape` is `None`. Otherwise, the multi-head outputs are
-        projected to the shape specified by `output_shape`. `cache` is the
-        updated cache.
+        One of the following:
+        - `attention_output`
+        - `(attention_output, attention_scores)`
+        - `(attention_output, cache)`
+        - `(attention_output, attention_scores, cache)`
     """
 
     def call(
@@ -72,6 +73,7 @@ def call(
         attention_mask=None,
         cache=None,
         cache_update_index=None,
+        return_attention_scores=False,
         training=None,
     ):
         if key is None:
@@ -113,11 +115,17 @@ def call(
             key=key,
             value=value,
             attention_mask=attention_mask,
+            return_attention_scores=return_attention_scores,
             training=training,
         )
 
         attention_output = self._output_dense(attention_output)
 
-        if cache is not None:
+        if return_attention_scores and cache is not None:
+            return attention_output, attention_scores, cache
+        elif return_attention_scores:
+            return attention_output, attention_scores
+        elif cache is not None:
             return attention_output, cache
-        return attention_output
+        else:
+            return attention_output

keras_hub/src/layers/modeling/cached_multi_head_attention_test.py

@@ -8,6 +8,14 @@
 
 
 class CachedMultiHeadAttentionTest(TestCase):
+    def setUp(self):
+        super().setUp()
+        self.batch_size = 2
+        self.seq_len = 5
+        self.num_heads = 2
+        self.key_dim = 4
+        self.hidden_dim = self.num_heads * self.key_dim
+
     def test_layer_behaviors(self):
         self.run_layer_test(
             cls=CachedMultiHeadAttention,
@@ -26,20 +34,25 @@ def test_layer_behaviors(self):
         )
 
     def test_cache_call_is_correct(self):
-        batch_size = 2
-        seq_len = 5
-        num_heads = 2
-        key_dim = 4
-        hidden_dim = num_heads * key_dim
-
-        input_shape = (batch_size, seq_len, hidden_dim)
+        input_shape = (self.batch_size, self.seq_len, self.hidden_dim)
         x = random.uniform(shape=input_shape)
-        input_cache = ops.zeros((batch_size, 2, seq_len, num_heads, key_dim))
+        input_cache = ops.zeros(
+            (
+                self.batch_size,
+                2,
+                self.seq_len,
+                self.num_heads,
+                self.key_dim,
+            )
+        )
         # Use a causal mask.
-        mask = ops.tril(ops.ones((seq_len, seq_len)))
+        mask = ops.tril(ops.ones((self.seq_len, self.seq_len)))
         outputs = ops.zeros_like(x)
 
-        layer = CachedMultiHeadAttention(num_heads=num_heads, key_dim=key_dim)
+        layer = CachedMultiHeadAttention(
+            num_heads=self.num_heads,
+            key_dim=self.key_dim,
+        )
         no_loop_outputs, no_loop_cache = layer(
             x,
             x,
@@ -50,8 +63,12 @@ def test_cache_call_is_correct(self):
 
         def loop_body(i, outputs, cache):
             # Compute the rest tokens.
-            next_input = ops.slice(x, (0, i, 0), (batch_size, 1, hidden_dim))
-            next_mask = ops.slice(mask, (i, 0), (1, seq_len))
+            next_input = ops.slice(
+                x,
+                (0, i, 0),
+                (self.batch_size, 1, self.hidden_dim),
+            )
+            next_mask = ops.slice(mask, (i, 0), (1, self.seq_len))
             next_output, cache = layer(
                 query=next_input,
                 value=next_input,
@@ -64,7 +81,7 @@ def loop_body(i, outputs, cache):
 
         def call(outputs, cache):
             _, outputs, cache = ops.while_loop(
-                cond=lambda i, outputs, cache: i < seq_len,
+                cond=lambda i, outputs, cache: i < self.seq_len,
                 body=loop_body,
                 loop_vars=[0, outputs, cache],
             )
@@ -75,26 +92,107 @@ def call(outputs, cache):
         self.assertAllClose(output, no_loop_outputs)
         self.assertAllClose(output_cache, no_loop_cache)
 
-    def test_training_propagation(self):
-        batch_size = 2
-        seq_len = 5
-        num_heads = 2
-        key_dim = 4
-        hidden_dim = num_heads * key_dim
+    def test_return_attention_scores(self):
+        x = random.uniform(
+            shape=(self.batch_size, self.seq_len, self.hidden_dim)
+        )
+
+        layer = CachedMultiHeadAttention(
+            num_heads=self.num_heads,
+            key_dim=self.key_dim,
+        )
+        output, attention_scores = layer(
+            x,
+            x,
+            return_attention_scores=True,
+        )
+
+        self.assertIsNotNone(attention_scores)
+        self.assertAllEqual(
+            output.shape,
+            [self.batch_size, self.seq_len, self.hidden_dim],
+        )
+        self.assertAllEqual(
+            attention_scores.shape,
+            [
+                self.batch_size,
+                self.num_heads,
+                self.seq_len,
+                self.seq_len,
+            ],
+        )
+        score_sums = ops.sum(attention_scores, axis=-1)
+        self.assertAllClose(
+            score_sums,
+            ops.ones_like(score_sums),
+            atol=1e-5,
+        )
 
-        input_shape = (batch_size, seq_len, hidden_dim)
+    def test_return_attention_scores_with_cache(self):
+        x = random.uniform(
+            shape=(self.batch_size, self.seq_len, self.hidden_dim)
+        )
+        input_cache = ops.zeros(
+            (
+                self.batch_size,
+                2,
+                self.seq_len,
+                self.num_heads,
+                self.key_dim,
+            )
+        )
+        mask = ops.tril(ops.ones((self.seq_len, self.seq_len)))
+
+        layer = CachedMultiHeadAttention(
+            num_heads=self.num_heads,
+            key_dim=self.key_dim,
+        )
+        output, attention_scores, output_cache = layer(
+            x,
+            x,
+            cache=input_cache,
+            cache_update_index=0,
+            attention_mask=mask,
+            return_attention_scores=True,
+        )
+
+        self.assertIsNotNone(attention_scores)
+        self.assertAllEqual(
+            output.shape,
+            [self.batch_size, self.seq_len, self.hidden_dim],
+        )
+        self.assertAllEqual(
+            attention_scores.shape,
+            [
+                self.batch_size,
+                self.num_heads,
+                self.seq_len,
+                self.seq_len,
+            ],
+        )
+        self.assertAllEqual(output_cache.shape, input_cache.shape)
+
+    def test_training_propagation(self):
+        input_shape = (self.batch_size, self.seq_len, self.hidden_dim)
         x = random.uniform(shape=input_shape)
 
         layer = CachedMultiHeadAttention(
-            num_heads=num_heads,
-            key_dim=key_dim,
+            num_heads=self.num_heads,
+            key_dim=self.key_dim,
             dropout=0.99999,  # Zeros out the outputs after the dropout layer
         )
         outputs = layer(x, x, training=True)
 
         # Custom computation with dropout rate sets to about 1.0
         value = layer._value_dense(x)
-        attention_scores = ops.zeros((batch_size, num_heads, seq_len, seq_len))
+        attention_scores = ops.zeros(
+            (
+                self.batch_size,
+                self.num_heads,
+                self.seq_len,
+                self.seq_len,
+            )
+        )
         attention_output = ops.einsum(
             layer._combine_equation, attention_scores, value
         )