Add optional Gated Attention

keras/src/backend/jax/trainer.py

@@ -1101,9 +1101,9 @@ def _get_distributed_iterator(self, distribution):
         for data in self.data_adapter.get_jax_iterator():
             if layouts is None:
                 layouts = tree.map_structure(
-                    lambda d: distribution.get_data_layout(
-                        d.shape
-                    ).backend_layout,
+                    lambda d: (
+                        distribution.get_data_layout(d.shape).backend_layout
+                    ),
                     data,
                 )
             yield _distribute_data(data, layouts)

keras/src/backend/numpy/core.py

@@ -396,8 +396,8 @@ def while_loop(
     maximum_iterations=None,
 ):
     current_iter = 0
-    iteration_check = (
-        lambda iter: maximum_iterations is None or iter < maximum_iterations
+    iteration_check = lambda iter: (
+        maximum_iterations is None or iter < maximum_iterations
     )
     is_tuple = isinstance(loop_vars, (tuple, list))
     loop_vars = tuple(loop_vars) if is_tuple else (loop_vars,)

keras/src/backend/torch/core.py

@@ -650,8 +650,8 @@ def while_loop(
     maximum_iterations=None,
 ):
     current_iter = 0
-    iteration_check = (
-        lambda iter: maximum_iterations is None or iter < maximum_iterations
+    iteration_check = lambda iter: (
+        maximum_iterations is None or iter < maximum_iterations
     )
     is_tuple = isinstance(loop_vars, (tuple, list))
     loop_vars = tuple(loop_vars) if is_tuple else (loop_vars,)

keras/src/layers/attention/grouped_query_attention.py

@@ -49,6 +49,12 @@ class GroupedQueryAttention(Layer):
         activity_regularizer: Regularizer for dense layer activity.
         kernel_constraint: Constraint for dense layer kernels.
         bias_constraint: Constraint for dense layer kernels.
+        use_gate: Boolean, whether to apply a gated attention mechanism.
+            When True, an additional gating branch is added based on the
+            (NeurIPS 2025 Best Paper)[https://arxiv.org/abs/2505.06708].
+            It applies a sigmoid-activated linear projection to the query
+            which then gates the attention output. This helps improve training
+            stability and eliminates "attention sinks".
         seed: Optional integer to seed the dropout layer.
 
     Call arguments:
@@ -102,6 +108,7 @@ def __init__(
         activity_regularizer=None,
         kernel_constraint=None,
         bias_constraint=None,
+        use_gate=False,
         seed=None,
         **kwargs,
     ):
@@ -117,6 +124,7 @@ def __init__(
         self.num_repeats = num_query_heads // num_key_value_heads
         self.dropout = dropout
         self.use_bias = use_bias
+        self.use_gate = use_gate
         self._flash_attention = flash_attention or is_flash_attention_enabled()
         self.kernel_initializer = initializers.get(kernel_initializer)
         self.bias_initializer = initializers.get(bias_initializer)
@@ -170,7 +178,16 @@ def build(
             **self._get_common_kwargs_for_sublayer(),
         )
         self._key_dense.build(key_shape)
-
+        if self.use_gate:
+            self._gate_dense = EinsumDense(
+                "bqm,muh->bquh",
+                output_shape=(None, self.num_query_heads, self.head_dim),
+                bias_axes="uh" if self.use_bias else None,
+                activation="sigmoid",
+                name="gate",
+                **self._get_common_kwargs_for_sublayer(),
+            )
+            self._gate_dense.build(key_shape)
         self._value_dense = EinsumDense(
             "bkm,mvh->bkvh",
             output_shape=(None, self.num_key_value_heads, self.head_dim),
@@ -247,7 +264,8 @@ def call(
             attention_mask=attention_mask,
             use_causal_mask=use_causal_mask,
         )
-
+        if self.use_gate:
+            gate = self._gate_dense(query)
         query = self._query_dense(query)
         key = self._key_dense(key)
         value = self._value_dense(value)
@@ -266,10 +284,11 @@ def call(
             attention_mask=attention_mask,
             training=training,
         )
-
-        output = self._output_dense(
-            output
-        )  # (batch_dim, target_seq_len, feature_dim)
+        # (batch_dim, target_seq_len, feature_dim)
+        if self.use_gate:
+            output = self._output_dense(gate * output)
+        else:
+            output = self._output_dense(output)
 
         if return_attention_scores:
             return output, scores
@@ -483,6 +502,7 @@ def get_config(self):
             "num_query_heads": self.num_query_heads,
             "num_key_value_heads": self.num_key_value_heads,
             "use_bias": self.use_bias,
+            "use_gate": self.use_gate,
             "dropout": self.dropout,
             "kernel_initializer": initializers.serialize(
                 self.kernel_initializer

keras/src/layers/attention/grouped_query_attention_test.py

@@ -60,6 +60,44 @@ def test_basics(self):
             run_training_check=False,
         )
 
+        self.run_layer_test(
+            layers.GroupedQueryAttention,
+            init_kwargs={
+                "num_query_heads": 2,
+                "num_key_value_heads": 2,
+                "head_dim": 2,
+                "use_gate": True,
+            },
+            input_shape={"query_shape": (2, 8, 16), "value_shape": (2, 4, 16)},
+            expected_output_shape=(2, 8, 16),
+            expected_num_trainable_weights=10,
+            expected_num_non_trainable_weights=0,
+            expected_num_seed_generators=0,
+            expected_num_losses=0,
+            supports_masking=True,
+            run_training_check=False,
+        )
+
+        self.run_layer_test(
+            layers.GroupedQueryAttention,
+            init_kwargs={
+                "num_query_heads": 2,
+                "num_key_value_heads": 2,
+                "head_dim": 2,
+                "use_bias": False,
+                "dropout": 0.5,
+                "use_gate": True,
+            },
+            input_shape={"query_shape": (2, 8, 16), "value_shape": (2, 4, 16)},
+            expected_output_shape=(2, 8, 16),
+            expected_num_trainable_weights=5,
+            expected_num_non_trainable_weights=0,
+            expected_num_seed_generators=1,
+            expected_num_losses=0,
+            supports_masking=True,
+            run_training_check=False,
+        )
+
     @pytest.mark.skipif(
         backend.backend() not in ("jax", "torch"),
         reason="Flash attention only supported on JAX and Torch",
@@ -187,6 +225,26 @@ def test_compute_output_shape(
         )
         self.assertEqual(output.shape, comp_output_shape)
 
+        layer = layers.GroupedQueryAttention(
+            num_query_heads=num_query_heads,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=2,
+            use_gate=True,
+        )
+        batch_size = 7
+        query_shape = (batch_size,) + query_dims
+        value_shape = (batch_size,) + value_dims
+        key_shape = (batch_size,) + key_dims if key_dims else None
+
+        query = np.ones(query_shape)
+        value = np.ones(value_shape)
+        key = np.ones(key_shape) if key_shape else None
+        output = layer(query=query, value=value, key=key)
+        comp_output_shape = layer.compute_output_shape(
+            query_shape, value_shape, key_shape
+        )
+        self.assertEqual(output.shape, comp_output_shape)
+
     @parameterized.named_parameters(
         ("query_value_dim_mismatch", (2, 4, 8), (2, 2, 7), 2),
         ("key_value_dim_mismatch", (2, 4, 8), (2, 2, 8), (2, 1, 7)),
@@ -207,6 +265,7 @@ def test_initializer(self):
             num_query_heads=16,
             num_key_value_heads=16,
             head_dim=64,
+            use_gate=True,
             kernel_initializer=initializers.TruncatedNormal(stddev=0.02),
         )
         layer.build((2, 4, 8), (2, 4, 8))
@@ -225,6 +284,11 @@ def test_initializer(self):
             layer._output_dense.kernel,
         )
 
+        self.assertNotAllClose(
+            layer._query_dense.kernel,
+            layer._gate_dense.kernel,
+        )
+
     @pytest.mark.skipif(
         backend.backend() == "numpy",
         reason="Numpy backend does not support masking.",
@@ -252,6 +316,30 @@ def test_query_mask_propagation(self):
                 )
         self.assertAllClose(masked_query._keras_mask, output._keras_mask)
 
+        try:
+            layer = layers.GroupedQueryAttention(
+                num_query_heads=2,
+                num_key_value_heads=2,
+                head_dim=2,
+                use_gate=True,
+            )
+            self.assertTrue(layer.supports_masking)
+            query = np.array(
+                [[1, 2, 3, 0, 0], [3, 3, 1, 1, 2], [1, 0, 0, 0, 0]]
+            )
+            masked_query = layers.Embedding(4, 8, mask_zero=True)(query)
+            value = np.random.normal(size=(3, 3, 8))
+            output = layer(query=masked_query, value=value)
+        except RuntimeError as e:
+            if e.args[0].startswith(
+                "(*bias): last dimension must be contiguous"
+            ):
+                self.skipTest(
+                    "PyTorch errors out on GPU: issue to track bug is here "
+                    "https://github.com/keras-team/keras/issues/20459"
+                )
+        self.assertAllClose(masked_query._keras_mask, output._keras_mask)
+
     @parameterized.named_parameters(("causal", True), ("not_causal", 0))
     @pytest.mark.skipif(
         backend.backend() == "numpy",
@@ -287,6 +375,34 @@ def test_masking(self, use_causal_mask):
         )
         self.assertAllClose(output, output_with_manual_mask)
 
+        layer = layers.GroupedQueryAttention(
+            num_query_heads=2, num_key_value_heads=2, head_dim=2, use_gate=True
+        )
+        query = np.array([[1, 2, 3, 0, 0], [3, 3, 1, 1, 2], [1, 0, 0, 0, 0]])
+        masked_query = layers.Embedding(4, 8, mask_zero=True)(query)
+        value = np.array([[5, 4, 0], [3, 0, 0], [2, 1, 1]])
+        masked_value = layers.Embedding(6, 8, mask_zero=True)(value)
+        output = layer(
+            query=masked_query,
+            value=masked_value,
+            use_causal_mask=use_causal_mask,
+        )
+        mask = np.array(
+            [[[1, 1, 0]] * 3 + [[0, 0, 0]] * 2]
+            + [[[1, 0, 0]] * 5]
+            + [[[1, 1, 1]] + [[0, 0, 0]] * 4]
+        ).astype(bool)
+        if use_causal_mask:
+            mask = mask & np.array(
+                [[[1, 0, 0], [1, 1, 0]] + [[1, 1, 1]] * 3]
+            ).astype(bool)
+        del masked_query._keras_mask
+        del masked_value._keras_mask
+        output_with_manual_mask = layer(
+            query=masked_query, value=masked_value, attention_mask=mask
+        )
+        self.assertAllClose(output, output_with_manual_mask)
+
     @parameterized.named_parameters(
         ("disable_flash_attention", False), ("enable_flash_attention", True)
     )

keras/src/layers/attention/multi_head_attention.py

@@ -64,6 +64,12 @@ class MultiHeadAttention(Layer):
         activity_regularizer: Regularizer for dense layer activity.
         kernel_constraint: Constraint for dense layer kernels.
         bias_constraint: Constraint for dense layer kernels.
+        use_gate: Boolean, whether to apply a gated attention mechanism.
+            When True, an additional gating branch is added based on the
+            (NeurIPS 2025 Best Paper)[https://arxiv.org/abs/2505.06708].
+            It applies a sigmoid-activated linear projection to the query
+            which then gates the attention output. This helps improve training
+            stability and eliminates "attention sinks".
         seed: Optional integer to seed the dropout layer.
 
     Call arguments:
@@ -117,6 +123,7 @@ def __init__(
         activity_regularizer=None,
         kernel_constraint=None,
         bias_constraint=None,
+        use_gate=False,
         seed=None,
         **kwargs,
     ):
@@ -127,6 +134,7 @@ def __init__(
         self._value_dim = value_dim if value_dim else key_dim
         self._dropout = dropout
         self._use_bias = use_bias
+        self._use_gate = use_gate
         if output_shape:
             if isinstance(output_shape, int):
                 output_shape = (output_shape,)
@@ -201,6 +209,7 @@ def get_config(self):
             "value_dim": self._value_dim,
             "dropout": self._dropout,
             "use_bias": self._use_bias,
+            "use_gate": self._use_gate,
             "output_shape": self._output_shape,
             "attention_axes": self._attention_axes,
             "kernel_initializer": initializers.serialize(
@@ -271,6 +280,18 @@ def build(
             **self._get_common_kwargs_for_sublayer(),
         )
         self._key_dense.build(key_shape)
+        if self._use_gate:
+            self._gate_dense = EinsumDense(
+                einsum_equation,
+                output_shape=_get_output_shape(
+                    output_rank - 1, [self._num_heads, self._key_dim]
+                ),
+                bias_axes=bias_axes if self._use_bias else None,
+                activation="sigmoid",
+                name="gate",
+                **self._get_common_kwargs_for_sublayer(),
+            )
+            self._gate_dense.build(key_shape)
         einsum_equation, bias_axes, output_rank = _build_proj_equation(
             value_rank - 1, bound_dims=1, output_dims=2
         )
@@ -549,6 +570,10 @@ def call(
         #   N = `num_attention_heads`
         #   H = `size_per_head`
 
+        # `gate` = [B, T, N, H]
+        if self._use_gate:
+            gate = self._gate_dense(query)
+
         # `query` = [B, T, N, H]
         query = self._query_dense(query)
 
@@ -565,7 +590,10 @@ def call(
             training,
             return_attention_scores,
         )
-        attention_output = self._output_dense(attention_output)
+        if self._use_gate:
+            attention_output = self._output_dense(attention_output * gate)
+        else:
+            attention_output = self._output_dense(attention_output)
 
         # Set mask on output if needed
         if query_mask is not None: