tf-models-nightly 2.18.0.dev20240917__py2.py3-none-any.whl → 2.18.0.dev20240919__py2.py3-none-any.whl

official/nlp/modeling/layers/transformer_encoder_block.py

@@ -115,6 +115,7 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
                use_sigmoid_attn=False,
                sigmoid_attn_bias=None,
                linformer_dim=None,
+               linformer_shared_kv_projection=True,
                **kwargs):
     """Initializes `TransformerEncoderBlock`.
 
@@ -194,6 +195,8 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         `block_sparse_attention.MultiHeadAttention`
       linformer_dim: Applies low-rank factorization on keys/values as in
         https://arxiv.org/pdf/2006.04768.
+      linformer_shared_kv_projection: If set, projection layer is shared for
+        keys and values.
       **kwargs: keyword arguments.
     """
     util.filter_kwargs(kwargs)
@@ -234,6 +237,7 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     self._use_sigmoid_attn = use_sigmoid_attn
     self._sigmoid_attn_bias = sigmoid_attn_bias
     self._linformer_dim = linformer_dim
+    self._linformer_shared_kv_projection = linformer_shared_kv_projection
     if self._num_kv_heads is not None and self._src_block_size is not None:
       raise ValueError(
           "Block sparse attention does not support Multi-query attention."
@@ -383,11 +387,13 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         dtype=tf.float32,
     )
     if self._linformer_dim is not None:
-      # Current implementation uses the same weights for keys and values.
-      # TODO(akandoor): Explore using different weights for keys and values.
+      if self._linformer_shared_kv_projection:
+        low_rank_dim = self._linformer_dim
+      else:
+        low_rank_dim = 2 * self._linformer_dim
       self._lowrank_kv_projection = tf_keras.layers.EinsumDense(
           "...bc,cd->...bd",
-          output_shape=(None, self._linformer_dim),
+          output_shape=(None, low_rank_dim),
           kernel_initializer=tf_utils.clone_initializer(
               self._kernel_initializer
           ),
@@ -444,6 +450,8 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         "tgt_block_size": self._tgt_block_size,
         "use_sigmoid_attn": self._use_sigmoid_attn,
         "sigmoid_attn_bias": self._sigmoid_attn_bias,
+        "linformer_dim": self._linformer_dim,
+        "linformer_shared_kv_projection": self._linformer_shared_kv_projection,
     }
     base_config = super().get_config()
     return dict(list(base_config.items()) + list(config.items()))
@@ -499,6 +507,8 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     if key_value is None:
       key_value = input_tensor
 
+    key = key_value
+    value = key_value
     if self._linformer_dim is not None:
       if attention_mask is not None:
         # Applying mask before the low rank factorization so that padding is
@@ -510,17 +520,28 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         attention_mask = None
       key_value = tf.transpose(key_value, [0, 2, 1])
       key_value = self._lowrank_kv_projection(key_value)
-      key_value = tf.transpose(key_value, [0, 2, 1])
-
+      if self._linformer_shared_kv_projection:
+        key_value = tf.transpose(key_value, [0, 2, 1])
+        key = key_value
+        value = key_value
+      else:
+        key = tf.transpose(key_value[:, :, :self._linformer_dim], [0, 2, 1])
+        value = tf.transpose(key_value[:, :, self._linformer_dim:], [0, 2, 1])
     if self._return_attention_scores:
       attention_output, attention_scores = self._attention_layer(
           query=target_tensor,
-          value=key_value,
+          key=key,
+          value=value,
           attention_mask=attention_mask,
-          return_attention_scores=True)
+          return_attention_scores=True,
+      )
     else:
       attention_output = self._attention_layer(
-          query=target_tensor, value=key_value, attention_mask=attention_mask)
+          query=target_tensor,
+          key=key,
+          value=value,
+          attention_mask=attention_mask,
+      )
     attention_output = self._attention_dropout(attention_output)
 
     if self._norm_first:

official/nlp/modeling/layers/transformer_encoder_block_test.py

@@ -800,7 +800,11 @@ class TransformerArgumentTest(tf.test.TestCase, parameterized.TestCase):
         output_tensor[1].shape.as_list(), expected_attention_scores_shape
     )
 
-  def test_low_rank_attention(self):
+  @parameterized.named_parameters(
+      ('unshared_kv_projection', False),
+      ('shared_kv_projection', True),
+  )
+  def test_low_rank_attention(self, shared_kv_projection):
     num_attention_heads = 8
     sequence_length = 21
     linformer_dim = 7
@@ -812,6 +816,7 @@ class TransformerArgumentTest(tf.test.TestCase, parameterized.TestCase):
         inner_activation='relu',
         return_attention_scores=True,
         linformer_dim=linformer_dim,
+        linformer_shared_kv_projection=shared_kv_projection,
     )
     # Create a 3-dimensional input (the first dimension is implicit).
     data_tensor = tf_keras.Input(shape=(sequence_length, width))

{tf_models_nightly-2.18.0.dev20240917.dist-info → tf_models_nightly-2.18.0.dev20240919.dist-info}/METADATA

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

{tf_models_nightly-2.18.0.dev20240917.dist-info → tf_models_nightly-2.18.0.dev20240919.dist-info}/RECORD

@@ -363,8 +363,8 @@ official/nlp/modeling/layers/tn_expand_condense_test.py,sha256=J52mXzoiuaXfR61kh
 official/nlp/modeling/layers/tn_transformer_expand_condense.py,sha256=gbGJOrgxJd1SyMGB6ME04FSxuZfHqsi94Xxt23l7368,11032
 official/nlp/modeling/layers/tn_transformer_test.py,sha256=Fh-EDRoAkhO7ccD3w3FsJHC51MnZySv8jBlHYnvKZMc,8893
 official/nlp/modeling/layers/transformer.py,sha256=yofIEOjZpcvDmHbcjBmkZrl5iSe6pLtMsetNbXmxDnY,20087
-official/nlp/modeling/layers/transformer_encoder_block.py,sha256=dxUCn9LckIJCpxJ8DRmiAU-4ycCmddXBTifLdihDmiU,24047
-official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=nbrfktOe0_WNhYYk0IlEJPf5d-9xtBoi2wDdO_FWF_k,32068
+official/nlp/modeling/layers/transformer_encoder_block.py,sha256=kiCQ4yGejmwRsJBKpmrwA1As4rFUekNYf9xGS052kyU,24766
+official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=cIunagl03W1tPkkt1BDVpGEpd-7ZwCqc3sPdzQOmpuc,32269
 official/nlp/modeling/layers/transformer_scaffold.py,sha256=m8TF4geBkm8-VJQiTpzMI6FSJZry6oa2vPO3FXCCClE,15704
 official/nlp/modeling/layers/transformer_scaffold_test.py,sha256=pqUGldhmAKROrd4eoCWmHNtKOdCO6PH_-EigcYnvIpE,19920
 official/nlp/modeling/layers/transformer_test.py,sha256=kC_9NcLbJnBbuTaE_7BW60EF8xG_QUoICj0t0gS7O4Q,5522
@@ -1222,9 +1222,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.18.0.dev20240917.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.18.0.dev20240917.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.18.0.dev20240917.dist-info/METADATA,sha256=l3aNnMlgUyV26Zw7dWix41njTiJ4a6o7gA-8SiI6Qq4,1432
-tf_models_nightly-2.18.0.dev20240917.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.18.0.dev20240917.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.18.0.dev20240917.dist-info/RECORD,,
+tf_models_nightly-2.18.0.dev20240919.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.18.0.dev20240919.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.18.0.dev20240919.dist-info/METADATA,sha256=z_bABqB1Cm6qU18Rku98ShkjheDYP-JbtDnzitB16cw,1432
+tf_models_nightly-2.18.0.dev20240919.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.18.0.dev20240919.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.18.0.dev20240919.dist-info/RECORD,,