[BUG] Wrong scale in `Attention2d` in `fused_attn` mode

[laclouis5]

Describe the bug
When the fused_attn is used, the scale of the attention is not specified in torch.nn.functional.scaled_dot_product_attention and the value defaults to q.size(-1) ** -0.5, which is different from the default from the Attention2d layer (num_heads ** -0.5).

This means that the results from the fused implementation and the vanilla one are different.

pytorch-image-models/timm/layers/attention2d.py

Lines 294 to 351 in dafe866

	class Attention2d(nn.Module):
	fused_attn: torch.jit.Final[bool]
	""" multi-head attention for 2D NCHW tensors"""
	def __init__(
	self,
	dim: int,
	dim_out: Optional[int] = None,
	num_heads: int = 32,
	bias: bool = True,
	expand_first: bool = False,
	head_first: bool = False,
	attn_drop: float = 0.,
	proj_drop: float = 0.
	):
	super().__init__()
	dim_out = dim_out or dim
	dim_attn = dim_out if expand_first else dim
	self.num_heads = num_heads
	self.dim_head = dim_attn // num_heads
	self.head_first = head_first
	self.scale = num_heads ** -0.5
	self.fused_attn = use_fused_attn()
	self.qkv = nn.Conv2d(dim, dim_attn * 3, 1, bias=bias)
	self.attn_drop = nn.Dropout(attn_drop)
	self.proj = nn.Conv2d(dim_attn, dim_out, 1, bias=bias)
	self.proj_drop = nn.Dropout(proj_drop)
	def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
	B, C, H, W = x.shape
	if self.head_first:
	q, k, v = self.qkv(x).view(B, self.num_heads, self.dim_head * 3, -1).chunk(3, dim=2)
	else:
	q, k, v = self.qkv(x).reshape(B, 3, self.num_heads, self.dim_head, -1).unbind(1)
	if self.fused_attn:
	x = torch.nn.functional.scaled_dot_product_attention(
	q.transpose(-1, -2).contiguous(),
	k.transpose(-1, -2).contiguous(),
	v.transpose(-1, -2).contiguous(),
	attn_mask=attn_mask,
	dropout_p=self.attn_drop.p if self.training else 0.,
	).transpose(-1, -2).reshape(B, -1, H, W)
	else:
	q = q * self.scale
	attn = q.transpose(-2, -1) @ k
	if attn_mask is not None:
	# NOTE: assumes mask is float and in correct shape
	attn = attn + attn_mask
	attn = attn.softmax(dim=-1)
	attn = self.attn_drop(attn)
	x = (v @ attn.transpose(-2, -1)).view(B, -1, H, W)
	x = self.proj(x)
	x = self.proj_drop(x)
	return x

Expected behavior
Same results for the two implementations.

Desktop (please complete the following information):

• OS: macOS
• This repository version: 1.0.12
• PyTorch version 2.5 (CPU)

Additional context
Add any other context about the problem here.

[rwightman]

@laclouis5 it's wrong but backwards from that, the .scale attr is wrong, looks like it might have been an autocomplete issue when I was switching the dim handling code / naming around to be consistent with other modules... should be self.scale = dim_head ** -0.5

This layer isn't actually active in any models right now, don't think I actually fully finished / tested it.

I have an impl in maxxvit that is working:

pytorch-image-models/timm/models/maxxvit.py

Lines 143 to 209 in e846b2c

	class Attention2d(nn.Module):
	fused_attn: Final[bool]
	""" multi-head attention for 2D NCHW tensors"""
	def __init__(
	self,
	dim: int,
	dim_out: Optional[int] = None,
	dim_head: int = 32,
	bias: bool = True,
	expand_first: bool = True,
	head_first: bool = True,
	rel_pos_cls: Callable = None,
	attn_drop: float = 0.,
	proj_drop: float = 0.
	):
	super().__init__()
	dim_out = dim_out or dim
	dim_attn = dim_out if expand_first else dim
	self.num_heads = dim_attn // dim_head
	self.dim_head = dim_head
	self.head_first = head_first
	self.scale = dim_head ** -0.5
	self.fused_attn = use_fused_attn()
	self.qkv = nn.Conv2d(dim, dim_attn * 3, 1, bias=bias)
	self.rel_pos = rel_pos_cls(num_heads=self.num_heads) if rel_pos_cls else None
	self.attn_drop = nn.Dropout(attn_drop)
	self.proj = nn.Conv2d(dim_attn, dim_out, 1, bias=bias)
	self.proj_drop = nn.Dropout(proj_drop)
	def forward(self, x, shared_rel_pos: Optional[torch.Tensor] = None):
	B, C, H, W = x.shape
	if self.head_first:
	q, k, v = self.qkv(x).view(B, self.num_heads, self.dim_head * 3, -1).chunk(3, dim=2)
	else:
	q, k, v = self.qkv(x).reshape(B, 3, self.num_heads, self.dim_head, -1).unbind(1)
	if self.fused_attn:
	attn_bias = None
	if self.rel_pos is not None:
	attn_bias = self.rel_pos.get_bias()
	elif shared_rel_pos is not None:
	attn_bias = shared_rel_pos
	x = torch.nn.functional.scaled_dot_product_attention(
	q.transpose(-1, -2).contiguous(),
	k.transpose(-1, -2).contiguous(),
	v.transpose(-1, -2).contiguous(),
	attn_mask=attn_bias,
	dropout_p=self.attn_drop.p if self.training else 0.,
	).transpose(-1, -2).reshape(B, -1, H, W)
	else:
	q = q * self.scale
	attn = q.transpose(-2, -1) @ k
	if self.rel_pos is not None:
	attn = self.rel_pos(attn)
	elif shared_rel_pos is not None:
	attn = attn + shared_rel_pos
	attn = attn.softmax(dim=-1)
	attn = self.attn_drop(attn)
	x = (v @ attn.transpose(-2, -1)).view(B, -1, H, W)
	x = self.proj(x)
	x = self.proj_drop(x)
	return x

Though should fix this one up too.. PR welcome or will add to the list.

[rwightman]

In that file MultiQueryAttention2d is well tested and used in MobileNetV4 hybrids.

The other module, MultiQueryAttentionV2 is also not used or well tested.

[laclouis5]

Great that is was not used, it can be fixed without much impact. I proposed a fix in #2387, which correct the scale for the unfused attention.