Add source: tmt/model/mesh.py

tmt/model/mesh.py

@@ -0,0 +1,89 @@
+"""
+mesh.py — MeshBuilder: constructs a dynamic token graph each forward pass.
+
+Novel vs standard: unlike Graph Transformers that use fixed pre-defined graphs,
+MeshBuilder recomputes the graph topology at every forward pass using cosine
+similarity of the current token representations.  Only the top-k nearest
+neighbours are connected, giving a sparse O(S·k) edge set instead of O(S²).
+"""
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+def build_mesh(
+    x: Tensor,
+    k: int,
+    batch_size: int,
+    seq_len: int,
+) -> Tuple[Tensor, Tensor]:
+    """
+    Build a dynamic kNN token graph from token embeddings.
+
+    Args:
+        x:          (B*S, D) flattened token representations
+        k:          number of nearest neighbours per token (graph_k)
+        batch_size: B
+        seq_len:    S
+
+    Returns:
+        edge_index: (2, E) COO edge list in torch_geometric format.
+                    Edges are within-batch — source and target indices are
+                    global node indices (0 … B*S-1).
+        edge_weight:(E,) cosine similarity of each edge.
+    """
+    N = batch_size * seq_len  # total nodes
+
+    # Normalise for cosine similarity
+    x_norm = F.normalize(x, p=2, dim=-1)  # (N, D)
+
+    # Block-diagonal cosine similarity — only connect tokens within same sample
+    # so information never leaks across batch items
+    sim_rows, sim_cols, sim_vals = [], [], []
+
+    for b in range(batch_size):
+        start = b * seq_len
+        end = start + seq_len
+        x_b = x_norm[start:end]  # (S, D)
+        sim = x_b @ x_b.T  # (S, S) cosine sim matrix
+
+        # Zero out self-connections
+        sim.fill_diagonal_(float("-inf"))
+
+        # Top-k neighbours per token
+        actual_k = min(k, seq_len - 1)
+        topk_vals, topk_idx = sim.topk(actual_k, dim=-1)  # (S, k)
+
+        src = torch.arange(seq_len, device=x.device).unsqueeze(1).expand(-1, actual_k)
+        src = src.reshape(-1) + start
+        dst = topk_idx.reshape(-1) + start
+        vals = topk_vals.reshape(-1)
+
+        sim_rows.append(src)
+        sim_cols.append(dst)
+        sim_vals.append(vals)
+
+    edge_index = torch.stack(
+        [torch.cat(sim_rows), torch.cat(sim_cols)], dim=0
+    )  # (2, E)
+    edge_weight = torch.cat(sim_vals)  # (E,)
+
+    return edge_index, edge_weight
+
+
+class MeshBuilder(torch.nn.Module):
+    """Thin nn.Module wrapper around build_mesh so it shows in model.repr."""
+
+    def __init__(self, k: int) -> None:
+        super().__init__()
+        self.k = k
+
+    def forward(self, x: Tensor, batch_size: int, seq_len: int) -> Tuple[Tensor, Tensor]:
+        return build_mesh(x, self.k, batch_size, seq_len)
+
+    def __repr__(self) -> str:
+        return f"MeshBuilder(k={self.k}, params=0)"