No public description

PiperOrigin-RevId: 674517247

Author: tensorflower-gardener

official/nlp/modeling/layers/transformer_encoder_block.py

@@ -112,6 +112,8 @@ def __init__(self,
                num_kv_heads=None,
                src_block_size=None,
                tgt_block_size=None,
+               use_sigmoid_attn=False,
+               sigmoid_attn_bias=None,
                **kwargs):
     """Initializes `TransformerEncoderBlock`.
 
@@ -185,6 +187,10 @@ def __init__(self,
         `block_sparse_attention.MultiHeadAttention` for more details.
       tgt_block_size: Target block size. Refer to
         `block_sparse_attention.MultiHeadAttention` for more details.
+      use_sigmoid_attn: This param is only used in
+        `block_sparse_attention.MultiHeadAttention`
+      sigmoid_attn_bias: This param is only used in
+        `block_sparse_attention.MultiHeadAttention`
       **kwargs: keyword arguments.
     """
     util.filter_kwargs(kwargs)
@@ -222,6 +228,8 @@ def __init__(self,
     self._num_kv_heads = num_kv_heads
     self._src_block_size = src_block_size
     self._tgt_block_size = tgt_block_size
+    self._use_sigmoid_attn = use_sigmoid_attn
+    self._sigmoid_attn_bias = sigmoid_attn_bias
     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."
@@ -285,6 +293,8 @@ def build(self, input_shape):
       attention_layer_kwargs.update(
           src_block_size=self._src_block_size,
           tgt_block_size=self._tgt_block_size,
+          use_sigmoid_attn=self._use_sigmoid_attn,
+          sigmoid_attn_bias=self._sigmoid_attn_bias,
           name="block_sparse_attention",
       )
       attention_fn = block_sparse_attention.MultiHeadAttention
@@ -413,6 +423,8 @@ def get_config(self):
         "num_kv_heads": self._num_kv_heads,
         "src_block_size": self._src_block_size,
         "tgt_block_size": self._tgt_block_size,
+        "use_sigmoid_attn": self._use_sigmoid_attn,
+        "sigmoid_attn_bias": self._sigmoid_attn_bias,
     }
     base_config = super().get_config()
     return dict(list(base_config.items()) + list(config.items()))

official/nlp/modeling/layers/transformer_encoder_block_test.py

@@ -14,6 +14,8 @@
 
 """Tests for Keras-based transformer block layer."""
 
+import math
+
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf, tf_keras
@@ -751,7 +753,11 @@ def test_attention_with_kv_heads(self, num_kv_heads):
         output_tensor[1].shape.as_list(), expected_attention_scores_shape
     )
 
-  def test_block_sparse_attention(self):
+  @parameterized.named_parameters(
+      ('use_softmax_attn', False),
+      ('use_sigmoid_attn', True),
+  )
+  def test_block_sparse_attention(self, use_sigmoid_attn):
     num_attention_heads = 8
     sequence_length = 21
     width = 80
@@ -765,6 +771,10 @@ def test_block_sparse_attention(self):
         return_attention_scores=True,
         src_block_size=src_block_size,
         tgt_block_size=tgt_block_size,
+        use_sigmoid_attn=use_sigmoid_attn,
+        sigmoid_attn_bias=-math.log(sequence_length)
+        if use_sigmoid_attn
+        else None,
     )
     # Create a 3-dimensional input (the first dimension is implicit).
     data_tensor = tf_keras.Input(shape=(sequence_length, width))