tf-models-nightly 2.17.0.dev20240604__py2.py3-none-any.whl → 2.17.0.dev20240606__py2.py3-none-any.whl

official/nlp/modeling/layers/multi_query_attention.py

@@ -0,0 +1,208 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+
+"""Keras-based attention layers to support multi-query attention.
+
+Based on https://arxiv.org/pdf/1911.02150.pdf and
+https://arxiv.org/pdf/2305.13245.pdf.
+"""
+
+import string
+from typing import Optional, Sequence, Union
+
+import tensorflow as tf, tf_keras
+
+_CHR_IDX = string.ascii_lowercase
+
+
+def _build_proj_equation(
+    free_dims: int, bound_dims: int, output_dims: int
+) -> ...:
+  """Builds an einsum equation for projections inside attention layer.
+
+  Args:
+    free_dims: The number of free dimensions which are copied from input to
+      output.
+    bound_dims: The number of bound dimensions part of input which are combined
+      with the kernel to produce output.
+    output_dims: The number of output dimensions.
+
+  Returns:
+    A tuple of einsum equation, bias axes and output rank.
+  """
+
+  input_str = ""
+  kernel_str = ""
+  output_str = ""
+  bias_axes = ""
+  letter_offset = 0
+  for i in range(free_dims):
+    char = _CHR_IDX[i + letter_offset]
+    input_str += char
+    output_str += char
+
+  letter_offset += free_dims
+  for i in range(bound_dims):
+    char = _CHR_IDX[i + letter_offset]
+    input_str += char
+    kernel_str += char
+
+  letter_offset += bound_dims
+  for i in range(output_dims):
+    char = _CHR_IDX[i + letter_offset]
+    kernel_str += char
+    output_str += char
+    bias_axes += char
+  equation = f"{input_str},{kernel_str}->{output_str}"
+
+  return equation, bias_axes, len(output_str)
+
+
+def _get_output_shape(
+    output_rank: int, known_last_dims: Sequence[int]
+) -> list[Optional[int]]:
+  return [None] * (output_rank - len(known_last_dims)) + list(known_last_dims)
+
+
+class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
+  """Multi-query attention layer."""
+
+  def __init__(self, num_kv_heads=None, **kwargs):
+    # num_kv_heads defines the number of key/value heads. A value of 1 means
+    # that the key/value heads are shared across all query heads. Any other
+    # value must be less than num_heads and must divide num_heads exactly. If
+    # num_kv_heads is greater than 1, query heads are split into groups of
+    # num_kv_heads.
+    super().__init__(**kwargs)
+    self._num_kv_heads = num_kv_heads or self._num_heads
+    assert (
+        self._num_kv_heads < self._num_heads
+    ), "num_kv_heads must be less than num_heads."
+    assert (
+        self._num_heads % self._num_kv_heads == 0
+    ), "num_kv_heads needs to divide num_heads exactly."
+
+  def _build_from_signature(
+      self,
+      query: Union[tf.Tensor, tf.TensorShape],
+      value: Union[tf.Tensor, tf.TensorShape],
+      key: Optional[Union[tf.Tensor, tf.TensorShape]] = None,
+  ):
+    """Builds layers and variables.
+
+    Once the method is called, self._built_from_signature will be set to
+    True.
+
+    Args:
+        query: Query tensor or TensorShape.
+        value: Value tensor or TensorShape.
+        key: Key tensor or TensorShape.
+    """
+    # pytype: disable=attribute-error
+    super()._build_from_signature(query=query, value=value, key=key)
+    # pytype: enable=attribute-error
+
+    with tf.init_scope():
+      # Key, value are shared across heads in multi-query attention.
+      # Overwrite the K, V projections, logits & attend einsum equations to
+      # remove the number of attention head dimension in K, V related tensors.
+      #
+      # The following capital letters are used to denote the tensor dimension
+      # parameters:
+      # B = batch size
+      # S = length of the key/value (source)
+      # T = length of the query (target)
+      # N = number of query attention heads
+      # K = number of key/value heads
+      # n = N // K
+      # H = dimensions of each attention head.
+      #
+      if self._num_kv_heads == 1:
+        output_dims = 1
+        key_last_dims = [self._key_dim]
+        value_last_dims = [self._value_dim]
+        self._dot_product_equation = "...SH,...TNH->...NTS"
+        self._combine_equation = "...NTS,...SH->...TNH"
+      else:
+        output_dims = 2
+        key_last_dims = [self._num_kv_heads, self._key_dim]
+        value_last_dims = [self._num_kv_heads, self._value_dim]
+        self._dot_product_equation = "...SKH,...TKnH->...nKTS"
+        self._combine_equation = "...nKTS,...SKH->...TnKH"
+
+      einsum_equation, bias_axes, output_rank = _build_proj_equation(
+          free_dims=self._key_shape.rank - 1,
+          bound_dims=1,
+          output_dims=output_dims,
+      )
+      self._key_dense = tf_keras.layers.EinsumDense(
+          einsum_equation,
+          output_shape=_get_output_shape(output_rank - 1, key_last_dims),
+          bias_axes=bias_axes if self._use_bias else None,
+          name="key",
+          **self._get_common_kwargs_for_sublayer(),
+      )
+      einsum_equation, bias_axes, output_rank = _build_proj_equation(
+          free_dims=self._value_shape.rank - 1,
+          bound_dims=1,
+          output_dims=output_dims,
+      )
+      self._value_dense = tf_keras.layers.EinsumDense(
+          einsum_equation,
+          output_shape=_get_output_shape(output_rank - 1, value_last_dims),
+          bias_axes=bias_axes if self._use_bias else None,
+          name="value",
+          **self._get_common_kwargs_for_sublayer(),
+      )
+
+  def _compute_attention(
+      self, query, key, value, attention_mask=None, training=None
+  ):
+    if self._num_kv_heads > 1:
+      query = tf.reshape(
+          query,
+          [
+              tf.shape(query)[0],
+              tf.shape(query)[1],
+              self._num_kv_heads,
+              self._num_heads // self._num_kv_heads,
+              tf.shape(query)[-1],
+          ],
+      )
+
+    # pytype: disable=attribute-error
+    attention_output, attention_scores = super()._compute_attention(
+        query, key, value, attention_mask=attention_mask, training=training
+    )
+    # pytype: enable=attribute-error
+    if self._num_kv_heads != 1:
+      attention_output = tf.reshape(
+          attention_output,
+          [
+              tf.shape(attention_output)[0],
+              tf.shape(attention_output)[1],
+              self._num_heads,
+              tf.shape(attention_output)[-1],
+          ],
+      )
+      attention_scores = tf.reshape(
+          attention_scores,
+          [
+              tf.shape(attention_scores)[0],
+              self._num_heads,
+              tf.shape(attention_scores)[-2],
+              tf.shape(attention_scores)[-1],
+          ],
+      )
+    return attention_output, attention_scores

official/nlp/modeling/layers/multi_query_attention_test.py

@@ -0,0 +1,215 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+
+"""Tests for multi-query attention layer."""
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf, tf_keras
+
+from official.nlp.modeling.layers import multi_query_attention
+
+
+class MultiQueryAttentionTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ("key_value_same_proj_mqa", 1, None, None, [40, 80]),
+      ("key_value_different_proj_mqa", 1, 32, 60, [40, 60]),
+      ("key_value_same_proj_gqa", 3, None, None, [40, 80]),
+      ("key_value_different_proj_gqa", 3, 32, 60, [40, 60]),
+  )
+  def test_non_masked_attention(
+      self, num_kv_heads, value_dim, output_shape, output_dims
+  ):
+    """Test that the attention layer can be created without a mask tensor."""
+    test_layer = multi_query_attention.MultiHeadAttention(
+        num_heads=12,
+        num_kv_heads=num_kv_heads,
+        key_dim=64,
+        value_dim=value_dim,
+        output_shape=output_shape,
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    query = tf_keras.Input(shape=(40, 80))
+    value = tf_keras.Input(shape=(20, 80))
+    output = test_layer(query=query, value=value)
+    self.assertEqual(output.shape.as_list(), [None] + output_dims)
+
+  @parameterized.named_parameters(
+      ("_mqa", 1),
+      ("_gqa", 3),
+  )
+  def test_non_masked_self_attention(self, num_kv_heads):
+    """Test with one input (self-attenntion) and no mask tensor."""
+    test_layer = multi_query_attention.MultiHeadAttention(
+        num_heads=12, num_kv_heads=num_kv_heads, key_dim=64
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    query = tf_keras.Input(shape=(40, 80))
+    output = test_layer(query, query)
+    self.assertEqual(output.shape.as_list(), [None, 40, 80])
+
+  @parameterized.named_parameters(
+      ("_mqa", 1),
+      ("_gqa", 3),
+  )
+  def test_attention_scores(self, num_kv_heads):
+    """Test attention outputs with coefficients."""
+    test_layer = multi_query_attention.MultiHeadAttention(
+        num_heads=12, num_kv_heads=num_kv_heads, key_dim=64
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    query = tf_keras.Input(shape=(40, 80))
+    output, coef = test_layer(query, query, return_attention_scores=True)
+    self.assertEqual(output.shape.as_list(), [None, 40, 80])
+    self.assertEqual(coef.shape.as_list(), [None, 12, 40, 40])
+
+  @parameterized.named_parameters(
+      ("_mqa", 1),
+      ("_gqa", 3),
+  )
+  def test_attention_scores_with_values(self, num_kv_heads):
+    """Test attention outputs with coefficients."""
+    test_layer = multi_query_attention.MultiHeadAttention(
+        num_heads=12, num_kv_heads=num_kv_heads, key_dim=64
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    query = tf_keras.Input(shape=(40, 80))
+    value = tf_keras.Input(shape=(60, 80))
+    output, coef = test_layer(query, value, return_attention_scores=True)
+    self.assertEqual(output.shape.as_list(), [None, 40, 80])
+    self.assertEqual(coef.shape.as_list(), [None, 12, 40, 60])
+
+  @parameterized.named_parameters(
+      ("with_bias_mqa", 1, True),
+      ("no_bias_mqa", 1, False),
+      ("with_bias_gqa", 2, True),
+      ("no_bias_gqa", 2, False),
+  )
+  def test_masked_attention(self, num_kv_heads, use_bias):
+    """Test with a mask tensor."""
+    test_layer = multi_query_attention.MultiHeadAttention(
+        num_heads=4, num_kv_heads=num_kv_heads, key_dim=2, use_bias=use_bias
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    batch_size = 3
+    query = tf_keras.Input(shape=(4, 8))
+    value = tf_keras.Input(shape=(2, 8))
+    mask_tensor = tf_keras.Input(shape=(4, 2))
+    output = test_layer(query=query, value=value, attention_mask=mask_tensor)
+
+    # Create a model containing the test layer.
+    model = tf_keras.Model([query, value, mask_tensor], output)
+
+    # Generate data for the input (non-mask) tensors.
+    from_data = 10 * np.random.random_sample((batch_size, 4, 8))
+    to_data = 10 * np.random.random_sample((batch_size, 2, 8))
+
+    # Invoke the data with a random set of mask data. This should mask at
+    # least one element.
+    mask_data = np.random.randint(2, size=(batch_size, 4, 2))
+    masked_output_data = model.predict([from_data, to_data, mask_data])
+
+    # Invoke the same data, but with a null mask (where no elements are
+    # masked).
+    null_mask_data = np.ones((batch_size, 4, 2))
+    unmasked_output_data = model.predict([from_data, to_data, null_mask_data])
+
+    # Because one data is masked and one is not, the outputs should not be
+    # the same.
+    self.assertNotAllClose(masked_output_data, unmasked_output_data)
+
+    # Tests the layer with three inputs: Q, K, V.
+    key = tf_keras.Input(shape=(2, 8))
+    output = test_layer(
+        query, value=value, key=key, attention_mask=mask_tensor
+    )
+    model = tf_keras.Model([query, value, key, mask_tensor], output)
+
+    masked_output_data = model.predict(
+        [from_data, to_data, to_data, mask_data]
+    )
+    unmasked_output_data = model.predict(
+        [from_data, to_data, to_data, null_mask_data]
+    )
+    # Because one data is masked and one is not, the outputs should not be
+    # the same.
+    self.assertNotAllClose(masked_output_data, unmasked_output_data)
+
+    if use_bias:
+      self.assertLen(test_layer._query_dense.trainable_variables, 2)
+      self.assertLen(test_layer._output_dense.trainable_variables, 2)
+    else:
+      self.assertLen(test_layer._query_dense.trainable_variables, 1)
+      self.assertLen(test_layer._output_dense.trainable_variables, 1)
+
+  @parameterized.named_parameters(
+      ("_mqa", 1),
+      ("_gqa", 2),
+  )
+  def test_masked_attention_with_scores(self, num_kv_heads):
+    """Test with a mask tensor."""
+    test_layer = multi_query_attention.MultiHeadAttention(
+        num_heads=4, num_kv_heads=num_kv_heads, key_dim=2
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    batch_size = 3
+    query = tf_keras.Input(shape=(4, 8))
+    value = tf_keras.Input(shape=(2, 8))
+    mask_tensor = tf_keras.Input(shape=(4, 2))
+    output = test_layer(query=query, value=value, attention_mask=mask_tensor)
+
+    # Create a model containing the test layer.
+    model = tf_keras.Model([query, value, mask_tensor], output)
+
+    # Generate data for the input (non-mask) tensors.
+    from_data = 10 * np.random.random_sample((batch_size, 4, 8))
+    to_data = 10 * np.random.random_sample((batch_size, 2, 8))
+
+    # Invoke the data with a random set of mask data. This should mask at
+    # least one element.
+    mask_data = np.random.randint(2, size=(batch_size, 4, 2))
+    masked_output_data = model.predict([from_data, to_data, mask_data])
+
+    # Invoke the same data, but with a null mask (where no elements are
+    # masked).
+    null_mask_data = np.ones((batch_size, 4, 2))
+    unmasked_output_data = model.predict([from_data, to_data, null_mask_data])
+
+    # Because one data is masked and one is not, the outputs should not be
+    # the same.
+    self.assertNotAllClose(masked_output_data, unmasked_output_data)
+
+    # Create a model containing attention scores.
+    output, scores = test_layer(
+        query=query,
+        value=value,
+        attention_mask=mask_tensor,
+        return_attention_scores=True,
+    )
+    model = tf_keras.Model([query, value, mask_tensor], [output, scores])
+    masked_output_data_score, masked_score = model.predict(
+        [from_data, to_data, mask_data]
+    )
+    unmasked_output_data_score, unmasked_score = model.predict(
+        [from_data, to_data, null_mask_data]
+    )
+    self.assertNotAllClose(masked_output_data_score, unmasked_output_data_score)
+    self.assertAllClose(masked_output_data, masked_output_data_score)
+    self.assertAllClose(unmasked_output_data, unmasked_output_data_score)
+    self.assertNotAllClose(masked_score, unmasked_score)
+
+
+if __name__ == "__main__":
+  tf.test.main()

{tf_models_nightly-2.17.0.dev20240604.dist-info → tf_models_nightly-2.17.0.dev20240606.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.17.0.dev20240604
+Version: 2.17.0.dev20240606
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.17.0.dev20240604.dist-info → tf_models_nightly-2.17.0.dev20240606.dist-info}/RECORD

@@ -329,6 +329,8 @@ official/nlp/modeling/layers/moe.py,sha256=azoK1Cz5l-5yT-FOtb7VYSbePTxmoje-2XAcq
 official/nlp/modeling/layers/moe_test.py,sha256=sr4hys2su_fcf9-6t0awqOkico92pQ3I2ApVF2rvvf0,9414
 official/nlp/modeling/layers/multi_channel_attention.py,sha256=fInxZUAPrrgCoCkM62JylM_vf5Wp8D0uZAvd5yuooVU,7322
 official/nlp/modeling/layers/multi_channel_attention_test.py,sha256=KcKuq22gOOq8HBzxRs3gGCwDRhmJPzYxfzfnqorQ1fw,1922
+official/nlp/modeling/layers/multi_query_attention.py,sha256=fFPBa9IBVj_O5x5OfGuHUFnJmiouNL6F1KsBCeHUqwM,6978
+official/nlp/modeling/layers/multi_query_attention_test.py,sha256=3VFF2hz85YExWPwdbhYWaSrIaSOkC1x7axdGfXr0W90,8512
 official/nlp/modeling/layers/on_device_embedding.py,sha256=FgsHyRXf5TWVTyo4OeKImmrTnn4uOPJgS3AGKzKMWYY,4582
 official/nlp/modeling/layers/on_device_embedding_test.py,sha256=M-LUba4QXV37s9Cx7aH8LL3bz_YotC6qITmWRI7Fhjk,8589
 official/nlp/modeling/layers/pack_optimization.py,sha256=C2prsYZMSkL8FBjz6Syc_Tu4JgzppaeIHyGDDoWzs8c,10289
@@ -1208,9 +1210,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.17.0.dev20240604.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.17.0.dev20240604.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.17.0.dev20240604.dist-info/METADATA,sha256=In6-OiPJgknyfwf_yWw59CIEXWMCleXsPf7P8qKoxZM,1432
-tf_models_nightly-2.17.0.dev20240604.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.17.0.dev20240604.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.17.0.dev20240604.dist-info/RECORD,,
+tf_models_nightly-2.17.0.dev20240606.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240606.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240606.dist-info/METADATA,sha256=lAg2jWgOJTr2BqPQAYugSYGcJOYFHeKst5MSY8bdGog,1432
+tf_models_nightly-2.17.0.dev20240606.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240606.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240606.dist-info/RECORD,,