fix overflow in softmax in attention

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "jaxtyping>=0.2.38",
     "torch>=2.5.1",
     "setuptools>=78.1.0",
+    "einops>=0.8.1",
 ]
 
 [build-system]
@@ -26,6 +27,7 @@ dev = [
     "pytest>=8.3.4",
     "sphinx",
     "ruff",
+    "ipykernel>=6.30.1",
 ]
 
 [tool.ruff.lint]
@@ -38,3 +40,6 @@ source = "vcs"
 version_scheme = "no-guess-dev"
 fallback_version = "0.0.0"
 local_scheme = "no-local-version" # see https://github.com/pypa/setuptools-scm/issues/455
+
+[[tool.uv.index]]
+url = "https://pypi.org/simple"

src/e3tools/__init__.py

@@ -1,5 +1,12 @@
-from ._scatter import scatter
+from ._scatter import scatter, scatter_softmax
 from ._radius import radius, radius_graph
 from ._pack_unpack import pack_irreps, unpack_irreps
 
-__all__ = ["scatter", "radius", "radius_graph", "pack_irreps", "unpack_irreps"]
+__all__ = [
+    "pack_irreps",
+    "radius",
+    "radius_graph",
+    "scatter",
+    "scatter_softmax",
+    "unpack_irreps",
+]

src/e3tools/_scatter.py

@@ -44,3 +44,20 @@ def scatter(
         f"{out.ndim=}, {index.ndim=} {out_shape=}, {in_shape=}, {dim=}"
     )
     return torch.scatter_reduce(out, dim, index, src, reduce, include_self=False)
+
+
+def scatter_softmax(
+    src: torch.Tensor, index: torch.Tensor, dim: int, dim_size: int | None = None
+):
+    index = broadcast(index, src, dim)
+
+    max = scatter(src, index, dim, dim_size, reduce="amax")
+    max = torch.gather(input=max, dim=dim, index=index)
+
+    scores = (src - max).exp()
+    z = scatter(scores, index, dim, dim_size, reduce="sum")
+    z = torch.gather(input=z, dim=dim, index=index)
+
+    scores = scores / z
+
+    return scores

src/e3tools/nn/_transformer.py

@@ -5,7 +5,7 @@
 import e3nn.o3
 from torch import nn
 
-from e3tools import scatter
+from e3tools import scatter, scatter_softmax
 
 from ._conv import Conv
 from ._interaction import LinearSelfInteraction
@@ -124,16 +124,13 @@ def forward(self, node_attr, edge_index, edge_attr, edge_sh):
         v = self.h_v.apply_per_edge(node_attr[src], edge_attr, edge_sh)
 
         # compute softmax
-        exp = self.dot(q[dst], k).exp()
-        z = scatter(exp, dst, dim=0, dim_size=N, reduce="mean")
-        alpha = exp / z[dst]
+        alpha = self.dot(q[dst], k)
+        alpha = scatter_softmax(alpha, dst, dim=0, dim_size=N)
 
-        attn = alpha.relu().sqrt()
-
-        out = scatter(attn * v, dst, dim=0, dim_size=N, reduce="sum")
+        out = scatter(alpha * v, dst, dim=0, dim_size=N, reduce="sum")
 
         if self.return_attention:
-            return out, attn
+            return out, alpha
         else:
             return out
 
@@ -252,20 +249,15 @@ def forward(self, node_attr, edge_index, edge_attr, edge_sh):
         k = k.view(E, self.num_heads, -1)
         v = v.view(E, self.num_heads, -1)
 
-        # compute softmax
-        exp = self.dot(q[dst], k).exp()
-        z = scatter(exp, dst, dim=0, dim_size=N, reduce="mean")
-        alpha = exp / z[dst]
-
-        attn = alpha.relu().sqrt()
-
-        out = scatter(attn * v, dst, dim=0, dim_size=N, reduce="sum").view(N, -1)
+        alpha = self.dot(q[dst], k)
+        alpha = scatter_softmax(alpha, dst, dim=0, dim_size=N)
+        out = scatter(alpha * v, dst, dim=0, dim_size=N, reduce="sum").view(N, -1)
 
         # use linear layer to transform back into original irreps
         out = self.lin_out(out)
 
         if self.return_attention:
-            return out, attn
+            return out, alpha
         else:
             return out
 

tests/test_equivariance.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 import e3nn
+
 from e3tools.nn import (
     Attention,
     Conv,

tests/test_scatter.py

@@ -0,0 +1,43 @@
+import einops
+import torch
+
+from e3tools import scatter, scatter_softmax
+
+
+def test_scatter_basic():
+    B = 4
+    D = 7
+    x = torch.randn(B, D)
+    ref = torch.sum(x, dim=1)
+
+    index = einops.repeat(torch.arange(B), "b -> (b d)", d=D)
+
+    out = scatter(x.view(-1), index, dim=0, dim_size=B, reduce="sum")
+
+    torch.testing.assert_close(ref, out)
+
+
+def test_scatter_softmax():
+    B = 4
+    D = 7
+    x = torch.randn(B, D)
+    ref = torch.nn.functional.softmax(x, dim=1)
+
+    index = einops.repeat(torch.arange(B), "b -> (b d)", d=D)
+
+    out = scatter_softmax(x.view(-1), index, dim=0, dim_size=B)
+
+    torch.testing.assert_close(ref.view(-1), out)
+
+
+def test_scatter_softmax_overflow():
+    B = 4
+    D = 7
+    x = torch.randn(B, D) + 100.0
+    ref = torch.nn.functional.softmax(x, dim=1)
+
+    index = einops.repeat(torch.arange(B), "b -> (b d)", d=D)
+
+    out = scatter_softmax(x.view(-1), index, dim=0, dim_size=B)
+
+    torch.testing.assert_close(ref.view(-1), out)