tf-models-nightly 2.20.0.dev20250721__py2.py3-none-any.whl → 2.20.0.dev20250722__py2.py3-none-any.whl

official/nlp/modeling/layers/transformer_encoder_block.py

@@ -82,41 +82,44 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
      Understanding](https://arxiv.org/abs/1810.04805)
   """
 
-  def __init__(self,
-               num_attention_heads,
-               inner_dim,
-               inner_activation,
-               output_range=None,
-               kernel_initializer="glorot_uniform",
-               bias_initializer="zeros",
-               kernel_regularizer=None,
-               bias_regularizer=None,
-               activity_regularizer=None,
-               kernel_constraint=None,
-               bias_constraint=None,
-               use_bias=True,
-               norm_first=False,
-               norm_epsilon=1e-12,
-               use_rms_norm=False,
-               output_dropout=0.0,
-               attention_dropout=0.0,
-               inner_dropout=0.0,
-               attention_initializer=None,
-               attention_axes=None,
-               use_query_residual=True,
-               key_dim=None,
-               value_dim=None,
-               output_last_dim=None,
-               diff_q_kv_att_layer_norm=False,
-               return_attention_scores=False,
-               num_kv_heads=None,
-               src_block_size=None,
-               tgt_block_size=None,
-               use_sigmoid_attn=False,
-               sigmoid_attn_bias=None,
-               linformer_dim=None,
-               linformer_shared_kv_projection=True,
-               **kwargs):
+  def __init__(
+      self,
+      num_attention_heads,
+      inner_dim,
+      inner_activation,
+      output_range=None,
+      kernel_initializer="glorot_uniform",
+      bias_initializer="zeros",
+      kernel_regularizer=None,
+      bias_regularizer=None,
+      activity_regularizer=None,
+      kernel_constraint=None,
+      bias_constraint=None,
+      use_bias=True,
+      norm_first=False,
+      norm_epsilon=1e-12,
+      use_rms_norm=False,
+      output_dropout=0.0,
+      attention_dropout=0.0,
+      inner_dropout=0.0,
+      attention_initializer=None,
+      attention_axes=None,
+      use_query_residual=True,
+      key_dim=None,
+      value_dim=None,
+      output_last_dim=None,
+      diff_q_kv_att_layer_norm=False,
+      return_attention_scores=False,
+      num_kv_heads=None,
+      src_block_size=None,
+      tgt_block_size=None,
+      use_sigmoid_attn=False,
+      sigmoid_attn_bias=None,
+      linformer_dim=None,
+      linformer_shared_kv_projection=True,
+      lowrank_query_seq_proj_dim=None,
+      **kwargs,
+  ):
     """Initializes `TransformerEncoderBlock`.
 
     Note: If `output_last_dim` is used and `use_query_residual` is `True`, the
@@ -197,6 +200,8 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         https://arxiv.org/pdf/2006.04768.
       linformer_shared_kv_projection: If set, projection layer is shared for
         keys and values.
+      lowrank_query_seq_proj_dim: If set, applies a projection layer on query
+        sequence to the given dimension. go/constformer-doc
       **kwargs: keyword arguments.
     """
     util.filter_kwargs(kwargs)
@@ -238,6 +243,7 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     self._sigmoid_attn_bias = sigmoid_attn_bias
     self._linformer_dim = linformer_dim
     self._linformer_shared_kv_projection = linformer_shared_kv_projection
+    self._lowrank_query_seq_proj_dim = lowrank_query_seq_proj_dim
     if (
         self._src_block_size is not None
         and self._num_kv_heads is not None
@@ -410,6 +416,21 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
           name="lowrank_kv_projection",
           **common_kwargs,
       )
+    if self._lowrank_query_seq_proj_dim is not None:
+      self._lowrank_query_seq_projection = tf_keras.layers.EinsumDense(
+          # Squash the sequence-length dimension; keep embedding as is.
+          "...ij,ik->...kj",
+          output_shape=(
+              self._lowrank_query_seq_proj_dim,
+              hidden_size,
+          ),
+          kernel_initializer=tf_utils.clone_initializer(
+              self._kernel_initializer
+          ),
+          bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
+          name="constformer_projection",
+          **common_kwargs,
+      )
     super().build(input_shape)
 
   def get_config(self):
@@ -461,10 +482,66 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         "sigmoid_attn_bias": self._sigmoid_attn_bias,
         "linformer_dim": self._linformer_dim,
         "linformer_shared_kv_projection": self._linformer_shared_kv_projection,
+        "lowrank_query_seq_proj_dim": self._lowrank_query_seq_proj_dim,
     }
     base_config = super().get_config()
     return dict(list(base_config.items()) + list(config.items()))
 
+  def _apply_lowrank_query_projection(
+      self,
+      query: tf.Tensor,
+      attention_mask: tf.Tensor | None,
+  ):
+    """Applies constformer projection to the source tensor."""
+
+    # Don't project the source tensor if the `lowrank_query_seq_projection`
+    # (constformer) dimension is the same as the input
+    # sequence dimension.
+    if (
+        self._lowrank_query_seq_proj_dim is None
+        or query.shape[1] == self._lowrank_query_seq_proj_dim
+    ):
+      return query
+    # Don't overwrite the attention mask.
+    query = self._apply_query_mask(attention_mask, query)
+    dtype = query.dtype
+    query = self._lowrank_query_seq_projection(query)
+    query = tf.cast(query, dtype)
+    return query
+
+  def _apply_query_mask(
+      self,
+      attention_mask: tf.Tensor | None,
+      query: tf.Tensor,
+  ):
+    """Applying mask before the low rank factorization so that padding is accounted for.
+
+    Applies mask to query only if the dimension of query matches the mask. This
+    is to avoid the projection from happening multiple times while stacking
+    the transformer layers.
+
+    Args:
+      attention_mask: The attention_mask tensor.
+      query: The query tensor.
+
+    Returns:
+      query: The query tensor after applying the mask.
+    """
+    if attention_mask is None:
+      return query
+    if attention_mask.shape[1] != query.shape[1]:
+      # Skip the mask application for query.
+      logging.info(
+          "Skipping mask application on query. Shape mismatch: %s vs %s",
+          attention_mask.shape,
+          query.shape,
+      )
+      return query
+
+    query_mask = tf.cast(attention_mask[:, :, 0], dtype=query.dtype)
+    query = query * tf.expand_dims(query_mask, axis=-1)
+    return query
+
   def call(self, inputs: Any, output_range: Optional[tf.Tensor] = None) -> Any:
     """Transformer self-attention encoder block call.
 
@@ -499,6 +576,12 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     if output_range:
       if self._norm_first:
         source_tensor = input_tensor[:, 0:output_range, :]
+        if self._use_query_residual:
+          # `source_tensor` is only used for the residual connection.
+          source_tensor = self._apply_lowrank_query_projection(
+              source_tensor, attention_mask
+          )
+
         input_tensor = self._attention_layer_norm(input_tensor)
         if key_value is not None:
           key_value = self._attention_layer_norm_kv(key_value)
@@ -508,11 +591,21 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     else:
       if self._norm_first:
         source_tensor = input_tensor
+        if self._use_query_residual:
+          # `source_tensor` is only used for the residual connection.
+          source_tensor = self._apply_lowrank_query_projection(
+              source_tensor, attention_mask
+          )
         input_tensor = self._attention_layer_norm(input_tensor)
         if key_value is not None:
           key_value = self._attention_layer_norm_kv(key_value)
       target_tensor = input_tensor
 
+    # Project the query to the constformer dimension.
+    target_tensor = self._apply_lowrank_query_projection(
+        target_tensor, attention_mask
+    )
+
     if key_value is None:
       key_value = input_tensor
 
@@ -523,7 +616,8 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         # Applying mask before the low rank factorization so that padding is
         # accounted for.
         query_mask = tf.cast(attention_mask[:, :, 0], dtype=target_tensor.dtype)
-        target_tensor = target_tensor * tf.expand_dims(query_mask, axis=-1)
+        if self._lowrank_query_seq_proj_dim is None:
+          target_tensor = target_tensor * tf.expand_dims(query_mask, axis=-1)
         key_mask = tf.cast(attention_mask[:, 0, :], dtype=target_tensor.dtype)
         key_value = key_value * tf.expand_dims(key_mask, axis=-1)
         attention_mask = None
@@ -534,8 +628,9 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         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])
+        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,

official/nlp/modeling/layers/transformer_encoder_block_test.py

@@ -844,6 +844,90 @@ class TransformerArgumentTest(tf.test.TestCase, parameterized.TestCase):
         output_tensor[1].shape.as_list(), expected_attention_scores_shape
     )
 
+  def test_low_rank_attention_with_constformer(self):
+    num_attention_heads = 8
+    sequence_length = 21
+    linformer_dim = 7
+    lowrank_query_seq_proj_dim = 10
+    width = 80
+    shared_kv_projection = False
+
+    test_layer = TransformerEncoderBlock(
+        num_attention_heads=num_attention_heads,
+        inner_dim=2048,
+        inner_activation='relu',
+        return_attention_scores=True,
+        linformer_dim=linformer_dim,
+        linformer_shared_kv_projection=shared_kv_projection,
+        lowrank_query_seq_proj_dim=lowrank_query_seq_proj_dim,
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    data_tensor = tf_keras.Input(shape=(sequence_length, width))
+    output_tensor = test_layer(data_tensor)
+
+    # The output from constformer has bottlenecked sequence length.
+    expected_layer_output_shape = [None, lowrank_query_seq_proj_dim, width]
+    # Note that attentions scores with Constformer don't have same
+    # interpretation as the original attention scores, since the sequence
+    # length is squashed.
+    expected_attention_scores_shape = [
+        None,
+        num_attention_heads,
+        lowrank_query_seq_proj_dim,
+        linformer_dim,
+    ]
+
+    self.assertIsInstance(output_tensor, tuple)
+    self.assertLen(output_tensor, 2)
+    # First is the standard output.
+    self.assertEqual(
+        output_tensor[0].shape.as_list(), expected_layer_output_shape
+    )
+    # Second is the attention scores.
+    self.assertEqual(
+        output_tensor[1].shape.as_list(), expected_attention_scores_shape
+    )
+
+  def test_low_rank_attention_with_constformer_no_linformer(self):
+    num_attention_heads = 8
+    sequence_length = 21
+    lowrank_query_seq_proj_dim = 10
+    width = 80
+
+    test_layer = TransformerEncoderBlock(
+        num_attention_heads=num_attention_heads,
+        inner_dim=2048,
+        inner_activation='relu',
+        return_attention_scores=True,
+        lowrank_query_seq_proj_dim=lowrank_query_seq_proj_dim,
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    data_tensor = tf_keras.Input(shape=(sequence_length, width))
+    output_tensor = test_layer(data_tensor)
+
+    # The output from constformer has bottlenecked sequence length.
+    expected_layer_output_shape = [None, lowrank_query_seq_proj_dim, width]
+    # Note that attentions scores with Constformer don't have same
+    # interpretation as the original attention scores, since the sequence
+    # length is squashed.
+    expected_attention_scores_shape = [
+        None,
+        num_attention_heads,
+        lowrank_query_seq_proj_dim,
+        sequence_length,
+    ]
+
+    self.assertIsInstance(output_tensor, tuple)
+    self.assertLen(output_tensor, 2)
+    # First is the standard output.
+    self.assertEqual(
+        output_tensor[0].shape.as_list(), expected_layer_output_shape
+    )
+    # Second is the attention scores.
+    self.assertEqual(
+        output_tensor[1].shape.as_list(), expected_attention_scores_shape
+    )
+
 
 if __name__ == '__main__':
   tf.test.main()

{tf_models_nightly-2.20.0.dev20250721.dist-info → tf_models_nightly-2.20.0.dev20250722.dist-info}/METADATA

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

{tf_models_nightly-2.20.0.dev20250721.dist-info → tf_models_nightly-2.20.0.dev20250722.dist-info}/RECORD

@@ -363,8 +363,8 @@ official/nlp/modeling/layers/tn_expand_condense_test.py,sha256=QWq1dJqQUPe5n69K3
 official/nlp/modeling/layers/tn_transformer_expand_condense.py,sha256=omzTkCBEk2TOkHEYDEBwve6WsOitX7IIJHzeKXdqDq0,11012
 official/nlp/modeling/layers/tn_transformer_test.py,sha256=pSCONEZRI4J9_6QLTJ3g_ynUYLrRXsJ1c2YMSiOV_60,8893
 official/nlp/modeling/layers/transformer.py,sha256=VjUO-gVj_PnavbT_vSrg5NDKMr0SRSiqSg5ktd42m5M,20087
-official/nlp/modeling/layers/transformer_encoder_block.py,sha256=fsYdA40A5kh8KvrEUyBLmv8UkDkV3eLdQb9mleocoM0,24930
-official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=PA_XJ4epjJLcgXjgVRoWnOvqszn_a2RKfFmbo5ow724,32406
+official/nlp/modeling/layers/transformer_encoder_block.py,sha256=E-WeoxsjByL-lkPAyEjDbvt1_3ghcIpCXnoLSeCDKFQ,27953
+official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=Y1Byz7RgF6puMIU3WRyoWghp5a3rYemgTRFRDohfy2Q,35402
 official/nlp/modeling/layers/transformer_scaffold.py,sha256=qmzhCJvbbFVF9zDqnfO4Zs2JDXwKhK7iEBOhsU6-KpQ,15704
 official/nlp/modeling/layers/transformer_scaffold_test.py,sha256=dRJwesTBKm-mF5mDHrHfVpVNnxa-Wx-fj_4ZHDPTpE0,19920
 official/nlp/modeling/layers/transformer_test.py,sha256=-pk9cdz9UlMpCIkGRkCKsMmjdRGi0seySaaB_2dwmXw,5522
@@ -1248,9 +1248,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=yiAneltAW3NHSj3fUSvHNBjfq0MGZ
 tensorflow_models/nlp/__init__.py,sha256=8uQd4wI6Zc4IJMPjtQifMeWVbPFkTxqYh66wfivCOL4,807
 tensorflow_models/uplift/__init__.py,sha256=NzaweFf4ZmhRb2l_fuV6bP-2N8oSO3xu6xJqVb1UmpY,999
 tensorflow_models/vision/__init__.py,sha256=ks420Ooqzi0hU7HnQpM5rylLaE-YcJdJkBx_umVaXlE,833
-tf_models_nightly-2.20.0.dev20250721.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.20.0.dev20250721.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.20.0.dev20250721.dist-info/METADATA,sha256=xZHsVrpfd06nYob1JYmhxn2mNk9GqGntpwFHCUJIjqs,1432
-tf_models_nightly-2.20.0.dev20250721.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.20.0.dev20250721.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.20.0.dev20250721.dist-info/RECORD,,
+tf_models_nightly-2.20.0.dev20250722.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.20.0.dev20250722.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.20.0.dev20250722.dist-info/METADATA,sha256=g7qc5gIL8nhGzTxMTsPABwW7C6sJwZBvjKjgDtWNxok,1432
+tf_models_nightly-2.20.0.dev20250722.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.20.0.dev20250722.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.20.0.dev20250722.dist-info/RECORD,,