floydnet 0.1.1__tar.gz → 0.1.2__tar.gz

{floydnet-0.1.1 → floydnet-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: floydnet
-Version: 0.1.1
+Version: 0.1.2
 Summary: Floyd Multi-Head Attention: a drop-in variant of PyTorch MHA with module and function APIs
 Project-URL: Homepage, https://github.com/ocx-lab/FloydNet
 Project-URL: Repository, https://github.com/ocx-lab/FloydNet
@@ -231,15 +231,17 @@ Requires-Dist: ruff>=0.5; extra == 'dev'
 Description-Content-Type: text/markdown
 
 # FloydNet
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Python](https://img.shields.io/badge/Python-3.9%2B-blue)](https://www.python.org/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.1%2B-orange)](https://pytorch.org/)
 
 Official implementation of an ICLR paper (TODO: add paper title, authors, and links/arXiv).
 
 ![Figure Pivotal Attention Mechanism for 2-Floyd/3-Floyd.](misc/pivotalattn2&3.png)
 
 This repository serves two audiences:
-
-- **Engineering users**: reusable PyTorch components (functional attention APIs and Transformer-style blocks) under `src/`.
-- **Research users**: scripts/configs to reproduce paper experiments under `example/`.
+- **Engineering users**: Reusable PyTorch components (functional attention APIs and Transformer-style blocks) under `src/`.
+- **Research users**: Scripts/configs to reproduce paper experiments (TSP, Graph Isomorphism, BREC) under `example/`.
 
 ---
 
@@ -267,10 +269,6 @@ For algorithmic details, hyperparameter choices, and analysis, please refer to t
 
 ---
 
-## Using the Attention / Transformer Block
-
-This section targets **engineering users** who want to import FloydNet as a dependency.
-
 ### Installation
 
 #### Option A: Install from PyPI

{floydnet-0.1.1 → floydnet-0.1.2}/README.md

@@ -1,13 +1,15 @@
 # FloydNet
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Python](https://img.shields.io/badge/Python-3.9%2B-blue)](https://www.python.org/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.1%2B-orange)](https://pytorch.org/)
 
 Official implementation of an ICLR paper (TODO: add paper title, authors, and links/arXiv).
 
 ![Figure Pivotal Attention Mechanism for 2-Floyd/3-Floyd.](misc/pivotalattn2&3.png)
 
 This repository serves two audiences:
-
-- **Engineering users**: reusable PyTorch components (functional attention APIs and Transformer-style blocks) under `src/`.
-- **Research users**: scripts/configs to reproduce paper experiments under `example/`.
+- **Engineering users**: Reusable PyTorch components (functional attention APIs and Transformer-style blocks) under `src/`.
+- **Research users**: Scripts/configs to reproduce paper experiments (TSP, Graph Isomorphism, BREC) under `example/`.
 
 ---
 
@@ -35,10 +37,6 @@ For algorithmic details, hyperparameter choices, and analysis, please refer to t
 
 ---
 
-## Using the Attention / Transformer Block
-
-This section targets **engineering users** who want to import FloydNet as a dependency.
-
 ### Installation
 
 #### Option A: Install from PyPI

{floydnet-0.1.1 → floydnet-0.1.2}/example/README.md

@@ -6,6 +6,14 @@ The paper reports results on **three benchmarks**:
 - BREC
 - TSP
 
+## 🚀 Key Results
+
+| Domain | Benchmark | Result |
+| :--- | :--- | :--- |
+| **Algorithmic** | CLRS-30 | **96.64%** (SOTA), effectively solving graph & string algorithms. |
+| **Optimization** | Non-Metric TSP | **88.3%** optimality on unseen sizes ($N=200$), vs 1.3% for Linkern heuristic. |
+| **Expressiveness** | Substructure Count | Near-zero error (MAE **0.001**) on complex substructure counting. |
+
 ### Graph Count
 
 The Graph Count benchmark and dataset construction follow:

{floydnet-0.1.1 → floydnet-0.1.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "floydnet"
-version = "0.1.1"
+version = "0.1.2"
 description = "Floyd Multi-Head Attention: a drop-in variant of PyTorch MHA with module and function APIs"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -50,7 +50,7 @@ include = [
   "LICENSE",
   "CITATION.cff",
   "CHANGELOG.md",
-  "paper/**",
+  "src/**",
 ]
 
 [tool.hatch.build.targets.wheel]

floydnet-0.1.2/src/floydnet/__init__.py

@@ -0,0 +1,8 @@
+from .functional import pivotal_attention, pivotal_attention3
+from .transformer import PivotalAttentionBlock
+
+__all__ = [
+    "pivotal_attention",
+    "pivotal_attention3",
+    "PivotalAttentionBlock",
+]

floydnet-0.1.2/src/floydnet/functional.py

@@ -0,0 +1,150 @@
+# Copyright 2025 Beijing Academy of Artificial Intelligence (BAAI)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+from typing import Optional
+import math
+
+import torch
+import torch.nn.functional as F
+
+def pivotal_attention(
+    q_ik: torch.Tensor,
+    k_ij: torch.Tensor,
+    k_jk: torch.Tensor,
+    v_ij: torch.Tensor,
+    v_jk: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout: float = 0.0,
+    scale: Optional[float] = None,
+    inf: float = 1e9,
+) -> torch.Tensor:
+    """Pivotal attention as described in "FLOYDNET: A LEARNING PARADIGM FOR GLOBAL RELATIONAL REASONING".
+
+    Shapes:
+        q_ik: (B, H, L_i, L_k, D)
+        k_ij: (B, H, L_i, L_j, D)
+        k_jk: (B, H, L_j, L_k, D)
+        v_ij: (B, H, L_i, L_j, D)
+        v_jk: (B, H, L_j, L_k, D)
+        attn_mask (optional): broadcastable to (B, H, L_i, L_k, L_j)
+
+    Args:
+        attn_mask: Additive mask (float) or boolean mask. If boolean, masked positions are set to -inf.
+        dropout: Dropout probability applied to attention weights (only effective if > 0).
+        scale: Optional custom scaling factor. If None, defaults to 1/sqrt(2*D).
+        inf: Value to use for -infinity in masks.
+
+    Returns:
+        Tensor of shape (B, H, L_i, L_k, D)
+    """
+    assert all([t.dim() == 5 for t in [q_ik, k_ij, k_jk, v_ij, v_jk]]), "All inputs must be 5D tensors"
+    B, H, L_i, L_k, D = q_ik.shape
+    L_j = k_ij.shape[3]
+    assert k_ij.shape == v_ij.shape == (B, H, L_i, L_j, D), "k_ij and v_ij must have shape (B, H, L_i, L_j, D)"
+    assert k_jk.shape == v_jk.shape == (B, H, L_j, L_k, D), "k_jk and v_jk must have shape (B, H, L_j, L_k, D)"
+
+    if scale is None:
+        scale = 1.0 / math.sqrt(2.0 * D)
+    q_ik = q_ik * scale
+
+    # Compute attention scores over the pivot dimension j: (B, H, L_i, L_k, L_j)
+    attn_scores = torch.einsum("bhikd,bhijd->bhikj", q_ik, k_ij) \
+                + torch.einsum("bhikd,bhjkd->bhikj", q_ik, k_jk)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_scores = attn_scores.masked_fill(attn_mask, -inf)
+        else:
+            attn_scores = attn_scores + attn_mask
+
+    attn_weights = torch.softmax(attn_scores, dim=-1)
+
+    if dropout > 0.0:
+        attn_weights = F.dropout(attn_weights, p=dropout)
+
+    y = torch.einsum("bhikj,bhijd->bhikd", attn_weights, v_ij) \
+      + torch.einsum("bhikj,bhjkd->bhikd", attn_weights, v_jk)
+
+    return y
+
+def pivotal_attention3(
+    q_ijk: torch.Tensor,
+    k_pjk: torch.Tensor,
+    k_ipk: torch.Tensor,
+    k_ijp: torch.Tensor,
+    v_pjk: torch.Tensor,
+    v_ipk: torch.Tensor,
+    v_ijp: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout: float = 0.0,
+    scale: Optional[float] = None,
+    inf: float = 1e9,
+) -> torch.Tensor:
+    """3-Pivotal attention as described in "FLOYDNET: A LEARNING PARADIGM FOR GLOBAL RELATIONAL REASONING".
+
+    Shapes:
+        q_ijk: (B, H, L_i, L_j, L_k, D)
+        k_pjk: (B, H, L_p, L_j, L_k, D)
+        k_ipk: (B, H, L_i, L_p, L_k, D)
+        k_ijp: (B, H, L_i, L_j, L_p, D)
+        v_pjk: (B, H, L_p, L_j, L_k, D)
+        v_ipk: (B, H, L_i, L_p, L_k, D)
+        v_ijp: (B, H, L_i, L_j, L_p, D)
+        attn_mask (optional): broadcastable to (B, H, L_i, L_j, L_k, L_p)
+
+    Args:
+        attn_mask: Additive mask (float) or boolean mask. If boolean, masked positions are set to -inf.
+        dropout: Dropout probability applied to attention weights (only effective if > 0).
+        scale: Optional custom scaling factor. If None, defaults to 1/sqrt(3*D).
+        inf: Value to use for -infinity in masks.
+
+    Returns:
+        Tensor of shape (B, H, L_i, l_j, L_k, D)
+    """
+    assert all([t.dim() == 6 for t in [q_ijk, k_pjk, k_ipk, k_ijp, v_pjk, v_ipk, v_ijp]]), "All inputs must be 6D tensors"
+    B, H, L_i, L_j, L_k, D = q_ijk.shape
+    L_p = k_pjk.shape[2]
+    assert k_pjk.shape == v_pjk.shape == (B, H, L_p, L_j, L_k, D), "k_pjk and v_pjk must have shape (B, H, L_p, L_j, L_k, D)"
+    assert k_ipk.shape == v_ipk.shape == (B, H, L_i, L_p, L_k, D), "k_ipk and v_ipk must have shape (B, H, L_i, L_p, L_k, D)"
+    assert k_ijp.shape == v_ijp.shape == (B, H, L_i, L_j, L_p, D), "k_ijp and v_ijp must have shape (B, H, L_i, L_j, L_p, D)"
+    
+    if scale is None:
+        scale = 1.0 / math.sqrt(3.0 * D)
+    q_ijk = q_ijk * scale
+
+    # Compute attention scores over the pivot dimension j: (B, H, L_i, L_j, L_k, L_p)
+    attn_scores = torch.einsum("bhijkd,bhpjkd->bhijkp", q_ijk, k_pjk) \
+                + torch.einsum("bhijkd,bhipkd->bhijkp", q_ijk, k_ipk) \
+                + torch.einsum("bhijkd,bhijpd->bhijkp", q_ijk, k_ijp)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_scores = attn_scores.masked_fill(attn_mask, -inf)
+        else:
+            attn_scores = attn_scores + attn_mask
+
+    attn_weights = torch.softmax(attn_scores, dim=-1)
+
+    if dropout > 0.0:
+        attn_weights = F.dropout(attn_weights, p=dropout)
+
+    y = torch.einsum("bhijkp,bhpjkd->bhijkd", attn_weights, v_pjk) \
+      + torch.einsum("bhijkp,bhipkd->bhijkd", attn_weights, v_ipk) \
+      + torch.einsum("bhijkp,bhijpd->bhijkd", attn_weights, v_ijp)
+
+    return y
+

floydnet-0.1.2/src/floydnet/transformer.py

@@ -0,0 +1,219 @@
+# Copyright 2025 Beijing Academy of Artificial Intelligence (BAAI)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import copy
+from typing import Optional, Tuple, Union, Callable
+
+import torch
+from torch import nn
+from .functional import pivotal_attention
+
+
+class Affine(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones((c, )))
+        self.bias = nn.Parameter(torch.zeros((c, )))
+
+    def forward(self, x: torch.Tensor):
+        return x * self.weight + self.bias
+
+
+def create_norm(norm_fn: Union[str, Callable], embed_dim: int, eps: float = 1e-5, **kwargs) -> nn.Module:
+    """Create a normalization module from a name or nn.Module.
+
+    Args:
+        norm_fn: Name or an nn.Module instance/class.
+        embed_dim: Embedding dimension (features) used to construct the norm.
+        eps: Numerical epsilon passed to the normalization layer if applicable.
+        **kwargs: Extra keyword arguments forwarded to the normalization layer.
+
+    Returns:
+        An nn.Module normalization instance.
+    """
+    if isinstance(norm_fn, str):
+        if norm_fn.lower() in ["layernorm", "ln"]:
+            return nn.LayerNorm(embed_dim, eps=eps, **kwargs)
+        elif norm_fn.lower() in ["batchnorm", "bn"]:
+            return nn.BatchNorm1d(embed_dim, eps=eps, **kwargs)
+        elif norm_fn.lower() in ["rmsnorm", "rms"]:
+            return nn.RMSNorm(embed_dim, eps=eps, **kwargs)
+        elif norm_fn.lower() in ["affine"]:
+            return Affine(embed_dim)
+        elif norm_fn.lower() in ["none", "identity"]:
+            return nn.Identity()
+        else:
+            raise ValueError(f"Unsupported norm_fn string: {norm_fn}")
+    elif callable(norm_fn):
+        if isinstance(norm_fn, nn.Module):
+            # deepcopy to avoid shared parameters
+            return copy.deepcopy(norm_fn)
+        elif isinstance(norm_fn, type) and issubclass(norm_fn, nn.Module):
+            return norm_fn(embed_dim, eps=eps, **kwargs)
+        else:
+            raise TypeError("norm_fn callable must be an nn.Module or nn.Module class")
+    else:
+        raise TypeError("norm_fn must be a string or callable")
+
+        
+def create_activation(activation_fn: Union[str, Callable]) -> nn.Module:
+    """Create an activation module from a name or nn.Module.
+
+    Args:
+        activation_fn: Name or an nn.Module instance/class.
+
+    Returns:
+        An nn.Module activation instance.
+    """
+    if isinstance(activation_fn, str):
+        if activation_fn.lower() == "relu":
+            return nn.ReLU()
+        elif activation_fn.lower() == "gelu":
+            return nn.GELU()
+        elif activation_fn.lower() == "silu":
+            return nn.SiLU()
+        else:
+            raise ValueError(f"Unsupported activation_fn string: {activation_fn}")
+    elif callable(activation_fn):
+        if isinstance(activation_fn, nn.Module):
+            return activation_fn
+        elif isinstance(activation_fn, type) and issubclass(activation_fn, nn.Module):
+            return activation_fn()
+        else:
+            raise TypeError("activation_fn callable must be an nn.Module or nn.Module class")
+    else:
+        raise TypeError("activation_fn must be a string or callable")
+
+
+class PivotalAttentionBlock(nn.Module):
+    """Transformer-style block that applies pivotal attention followed by an FFN.
+
+    Args:
+        embed_dim: Input/hidden embedding dimension (D).
+        num_heads: Number of attention heads (D must be divisible by num_heads).
+        dropout: Dropout probability for attention output and FFN output.
+        bias: Whether to include bias terms in linear layers.
+        ffn_expansion_ratio: Expansion ratio for the FFN hidden size.
+        norm_position: "pre" or "post" layer normalization placement.
+        activation_fn: Activation name/module used in the FFN.
+        norm_fn: Normalization name/module used in the block.
+    """
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        bias: bool = False,
+        ffn_expansion_ratio: int = 4,
+        norm_position: str = "pre",
+        activation_fn: Union[str, Callable] = "relu",
+        norm_fn: Union[str, Callable] = "layernorm",
+        enable_symmetric_mix: bool = True,
+        enable_ffn: bool = True,
+    ) -> None:
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.dropout = dropout
+        self.norm_position = norm_position.lower()
+        self.enable_ffn = enable_ffn
+        assert self.norm_position in ["pre", "post"], "norm_position must be 'pre' or 'post'"
+
+        self.enable_symmetric_mix = enable_symmetric_mix
+        if enable_symmetric_mix:
+            self.c_mix = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        self.c_qkv = nn.Linear(embed_dim, embed_dim * 5, bias=bias)
+        self.c_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.dropout_fn = nn.Dropout(dropout)
+        self.norm1 = create_norm(norm_fn, embed_dim)
+        if self.enable_ffn:
+            self.activation_fn = create_activation(activation_fn)
+            self.norm2 = create_norm(norm_fn, embed_dim)
+            self.ffn = nn.Sequential(
+                nn.Linear(embed_dim, ffn_expansion_ratio * embed_dim, bias=bias),
+                self.activation_fn,
+                nn.Linear(ffn_expansion_ratio * embed_dim, embed_dim, bias=bias),
+                nn.Dropout(dropout),
+            )
+            self.ffn_scale = nn.Parameter(torch.tensor(1.0, requires_grad=True))
+
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        """Initialize parameters using Xavier for projections and zeros for output heads."""
+        if self.enable_symmetric_mix:
+            nn.init.zeros_(self.c_mix.weight)
+        nn.init.xavier_uniform_(self.c_qkv.weight)
+        nn.init.zeros_(self.c_proj.weight)
+        if self.enable_ffn:
+            nn.init.xavier_uniform_(self.ffn[0].weight)
+            nn.init.zeros_(self.ffn[2].weight)
+        if self.c_qkv.bias is not None:
+            if self.enable_symmetric_mix:
+                nn.init.zeros_(self.c_mix.bias)
+            nn.init.zeros_(self.c_qkv.bias)
+            nn.init.zeros_(self.c_proj.bias)
+            if self.enable_ffn:
+                nn.init.zeros_(self.ffn[0].bias)
+                nn.init.zeros_(self.ffn[2].bias)
+
+    def attn(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor]) -> torch.Tensor:
+        """Apply pivotal attention over a (L x L) grid.
+
+        Args:
+            x: Input tensor of shape (B, L, L, D).
+            attn_mask: Optional mask broadcastable to (B, H, L, L, L).
+
+        Returns:
+            Tensor of shape (B, L, L, D) after attention projection and dropout.
+        """
+        B, L, _, D = x.shape
+        # [B, L, L, 5*D] -> 5 x [B, H, L, L, d]
+        qkv = torch.chunk(self.c_qkv(x), 5, dim=-1)
+        q_ik, k_ij, k_jk, v_ij, v_jk = map(
+            lambda t: t.view(B, L, L, self.num_heads, self.head_dim).permute(0, 3, 1, 2, 4),
+            qkv, 
+        )
+
+        # [B, H, L, L, d]
+        y = pivotal_attention(
+            q_ik, k_ij, k_jk, v_ij, v_jk,
+            attn_mask=attn_mask,
+            dropout=self.dropout if self.training else 0.0,
+        )
+        y = y.permute(0, 2, 3, 1, 4).contiguous().view(B, L, L, D)
+        y = self.c_proj(y)
+        y = self.dropout_fn(y)
+        return y
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.enable_symmetric_mix:
+            xT = self.c_mix(x.transpose(1, 2))
+        else:
+            xT = 0
+        if self.norm_position == "pre":
+            x = x + self.attn(self.norm1(x + xT), attn_mask)
+            if self.enable_ffn:
+                x = x + self.ffn(self.norm2(x)) * self.ffn_scale
+        else:
+            x = self.norm1(x + self.attn(x + xT, attn_mask))
+            if self.enable_ffn:
+                x = self.norm2(x + self.ffn(x)) * self.ffn_scale
+
+        return x