Simplify core attention return types to Tensor (#537)

Summary:

After removing all consumers of `head_mask`, `return_attn_weights`, and `attn_probs` in the previous commits, the core attention module can be simplified. This commit:
- Removes `head_mask` param from `scaled_dot_product_attention` and `SelfAttention.forward`
- Changes return types from `Tuple[Tensor, Tensor]` to `Tensor` (no longer returning attention probabilities)
- Removes `return_attn_weights` param and tuple unpacking logic from `MultiHeadAttention.forward`
- Cleans up unused imports (`Tuple`, `Union`)

No behavioral change — the attention computation itself is unchanged.

Reviewed By: OmarPavel

Differential Revision: D92927085

Author: stashuk-olek

tests/modules/layers/test_attention.py

@@ -244,7 +244,7 @@ def test_multi_head_attention_causal_use_cache(
 
 class TestScaledDotProductAttention:
     def test_scaled_dot_product_attention(self, q, kv):
-        output, weights = scaled_dot_product_attention(q, kv, kv)
+        output = scaled_dot_product_attention(q, kv, kv)
         actual = output
         expected = torch.tensor(
             [
@@ -263,31 +263,13 @@ def test_scaled_dot_product_attention(self, q, kv):
             ]
         )
         assert_expected(actual, expected, rtol=0, atol=1e-4)
-        actual = weights
-        expected = torch.tensor(
-            [
-                [
-                    [
-                        [
-                            [[0.8797, 0.1203], [0.5595, 0.4405]],
-                            [[0.0553, 0.9447], [0.4549, 0.5451]],
-                        ],
-                        [
-                            [[0.0419, 0.9581], [0.4391, 0.5609]],
-                            [[0.0297, 0.9703], [0.7313, 0.2687]],
-                        ],
-                    ]
-                ]
-            ]
-        )
-        assert_expected(actual, expected, rtol=0, atol=1e-4)
 
     def test_scaled_dot_product_attention_with_attention_mask(self, q, kv):
         attn_shape = torch.Size([1, 1, 2, 2, 2, 2])
         mask = torch.tensor([1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1]).view(
             attn_shape
         )
-        actual, _ = scaled_dot_product_attention(q, kv, kv, attention_mask=mask)
+        actual = scaled_dot_product_attention(q, kv, kv, attention_mask=mask)
         expected = torch.tensor(
             [
                 [
@@ -306,32 +288,8 @@ def test_scaled_dot_product_attention_with_attention_mask(self, q, kv):
         )
         assert_expected(actual, expected, rtol=0, atol=1e-4)
 
-    def test_scaled_dot_product_attention_with_head_mask(self, q, kv):
-        attn_shape = torch.Size([1, 1, 2, 2, 2, 2])
-        mask = torch.tensor([1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1]).view(
-            attn_shape
-        )
-        actual, _ = scaled_dot_product_attention(q, kv, kv, head_mask=mask)
-        expected = torch.tensor(
-            [
-                [
-                    [
-                        [
-                            [[-0.6195, 1.7705, 0.8023], [-0.2718, 1.2177, 1.4946]],
-                            [[-0.1561, 0.1993, 0.4882], [0.7800, -0.1800, 0.0093]],
-                        ],
-                        [
-                            [[0.2950, 1.2029, 1.7035], [0.1668, 0.7129, 1.0162]],
-                            [[0.0455, -0.0077, -0.0345], [0.7875, 0.0928, -0.7952]],
-                        ],
-                    ]
-                ]
-            ]
-        )
-        assert_expected(actual, expected, rtol=0, atol=1e-4)
-
     def test_scaled_dot_product_attention_with_dropout(self, q, kv):
-        actual, _ = scaled_dot_product_attention(q, kv, kv, attn_dropout=0.3)
+        actual = scaled_dot_product_attention(q, kv, kv, attn_dropout=0.3)
         expected = torch.tensor(
             [
                 [
@@ -365,7 +323,7 @@ def test_scaled_dot_product_attention_with_dropout(self, q, kv):
 
 def test_self_attention(self_attn, q, kv):
     k = v = kv
-    actual, _ = self_attn(q, k, v)
+    actual = self_attn(q, k, v)
     # Output of self attention should be same as scaled_dot_product_attention
     # since input dims are flattened
     expected = torch.tensor(

torchmultimodal/models/video_gpt/video_vqvae.py

@@ -66,7 +66,7 @@ def forward(
         old_shape = list(v.shape)
         v = v.flatten(end_dim=-3)
 
-        out, _ = scaled_dot_product_attention(
+        out = scaled_dot_product_attention(
             q,
             k,
             v,

torchmultimodal/modules/layers/attention.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional
 
 import torch
 from torch import nn, Tensor
@@ -29,8 +29,7 @@ def forward(
         k: Tensor,
         v: Tensor,
         attention_mask: Optional[Tensor] = None,
-        head_mask: Optional[Tensor] = None,
-    ) -> Tuple[Tensor, Tensor]:
+    ) -> Tensor:
         """
         Args:
             q (Tensor): Query input of shape ``(b, h, d1, ..., dn, dim_q)`` where ``h`` is the number of
@@ -42,11 +41,9 @@ def forward(
             attention_mask (Tensor, optional): Tensor of shape ``(b, h, q_dn, k_dn)`` where ``q_dn`` is the
                 dimension of the flattened query input along its latent dimensions and ``k_dn`` that of the
                 flattened key input. Contains 1s for positions to attend to and 0s for masked positions.
-            head_mask (Tensor, optional): Tensor of shape ``(b, h, q_dn, k_dn)``.
-                Contains 1s for positions to attend to and 0s for masked positions.
 
         Returns:
-            A tuple of output tensor and attention probabilities.
+            Output tensor.
         """
         _, _, *shape, _ = q.shape
 
@@ -55,16 +52,15 @@ def forward(
         k = k.flatten(start_dim=2, end_dim=-2)
         v = v.flatten(start_dim=2, end_dim=-2)
 
-        out, attn_probs = scaled_dot_product_attention(
+        out = scaled_dot_product_attention(
             q,
             k,
             v,
             attention_mask=attention_mask,
-            head_mask=head_mask,
             attn_dropout=self.attn_dropout if self.training else 0.0,
         )
 
-        return out.unflatten(2, shape), attn_probs
+        return out.unflatten(2, shape)
 
 
 class MultiHeadAttention(nn.Module):
@@ -121,11 +117,10 @@ def forward(
         self,
         q: Tensor,
         kv: Optional[Tensor] = None,
-        return_attn_weights: bool = False,
         use_cache: bool = False,
         causal: bool = False,
         **attn_kwargs: Any,
-    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    ) -> Tensor:
         """
         Args:
             q (Tensor): Query of shape ``(b, d1, ..., dn, dim_q)`` or ``(b, seq_len, dim_q)``
@@ -138,8 +133,7 @@ def forward(
             causal (bool): Whether to use causal attention or not. Default is ``False``.
 
         Returns:
-            * If ``return_attn_weights`` is ``True``: A tuple of output tensor and attention probabilities.
-            * If ``return_attn_weights`` is ``False``: A single output tensor.
+            Output tensor.
         """
         # If kv is specified use those inputs for cross-attention, otherwise use q
         k = v = q if kv is None else kv
@@ -169,51 +163,33 @@ def forward(
                 k, v = self.cache["k"], self.cache["v"]
 
         attn_out = self.attn(q, k, v, **attn_kwargs)
-        attn_probs = None
-        # Unpack if attn module also returns attn probs
-        if isinstance(attn_out, tuple):
-            attn_out, attn_probs = attn_out
         a = merge_multihead(attn_out)
         a = self.output(a)
 
-        if return_attn_weights:
-            return a, attn_probs
-        else:
-            return a
+        return a
 
 
 def scaled_dot_product_attention(
     q: Tensor,
     k: Tensor,
     v: Tensor,
     attention_mask: Optional[Tensor] = None,
-    head_mask: Optional[Tensor] = None,
     attn_dropout: float = 0.0,
-) -> Tuple[Tensor, Tensor]:
-    """Similar to PyTorch Core's _scaled_dot_product_attention but generalized
-    to handle n-dimensional input tokens (images, video) and support multihead.
-    Computes attention as described in Attention Is All You Need (Vaswani et al. 2017)
+) -> Tensor:
+    """Computes scaled dot-product attention. Similar to PyTorch Core's
+    ``scaled_dot_product_attention`` but generalized to handle n-dimensional
+    input tokens (images, video) and support multihead.
 
     Args:
-        q (Tensor): Query of shape ``(b, h, d1, ..., dn, dim_qk)`` or ``(b, h, seq_len, dim_qk)`` where
-            ``h`` is number of attention heads, ``d1, ..., dn`` are latent dimensions and ``dim_qk` is
-            the embedding dim of the query tensor.
-        k (Tensor): Key of shape ``(b, h, d1', ...., dn', dim_qk)`` or ``(b, h, seq_len', dim_qk)`` where
-            ``h`` is the number of attention heads, ``d1', ..., dn'` are latent dimensions and ``dim_qk``
-            is the key embedding dim aligned with query embedding dim,
-            see :class:`~torchmultimodal.modules.layers.attention.MultiHeadAttention`.
-        v (Tensor): Value of shape ``(b, h, d1', ..., dn', dim_v)`` or ``(b, h, seq_len', dim_v)`` where
-            ``h`` is the number of attention heads, ``d1', ..., dn'`` are latent dimensions and ``dim_v``
-            is the embedding dim of the value tensor.
+        q (Tensor): Query of shape ``(b, h, d1, ..., dn, dim_qk)`` or ``(b, h, seq_len, dim_qk)``.
+        k (Tensor): Key of shape ``(b, h, d1', ...., dn', dim_qk)`` or ``(b, h, seq_len', dim_qk)``.
+        v (Tensor): Value of shape ``(b, h, d1', ..., dn', dim_v)`` or ``(b, h, seq_len', dim_v)``.
         attention_mask (Tensor, optional): Tensor of shape ``(b, h, d1, ..., q_dn, k_dn)``.
-            Contains 1s for positions to attend to and 0s for masked positions. Applied before softmax.
-        head_mask (Tensor, optional): Tensor of shape ``(b, h, d1, ..., q_dn, k_dn)``.
             Contains 1s for positions to attend to and 0s for masked positions.
-            Applied after dropout, before matrix multiplication with values.
         attn_dropout (float): Probability of dropout after softmax. Default is ``0.0``.
 
     Returns:
-        A tuple of output tensor and attention probabilities.
+        Output tensor.
     """
 
     # Take the dot product between "query" and "key" and scale to get the raw attention scores.
@@ -232,13 +208,10 @@ def scaled_dot_product_attention(
     # This is actually dropping out entire tokens to attend to, which might
     # seem a bit unusual, but is taken from the original Transformer paper.
     attn = F.dropout(attn, p=attn_dropout)
-    # Mask heads if we want to
-    if head_mask is not None:
-        attn = attn * head_mask
     # For each query sum over the key/value dim with attention weights
     a = torch.matmul(attn, v)  # b, h, d1, ..., q_dn, c
 
-    return a, attn
+    return a
 
 
 def split_multihead(x: Tensor, n_head: int) -> Tensor: