1 files changed, 51 insertions, 0 deletions
diff --git a/worldalign/synth_fast_gate.py b/worldalign/synth_fast_gate.py
index 54b6891..01fb18f 100644
--- a/worldalign/synth_fast_gate.py
+++ b/worldalign/synth_fast_gate.py
@@ -343,3 +343,54 @@ def main() -> None:
 
 if __name__ == "__main__":
     main()
+
+
+def all_pair_swap_deltas(
+    permuted_text: torch.Tensor, visual: torch.Tensor
+) -> torch.Tensor:
+    """Energy change of every transposition at once, for the pairwise term.
+
+    The brute-force descent forms all N(N-1)/2 candidate permutations and
+    scores each through the full energy, which at 256 scenes means 32,640
+    matrices per step and, with the triangle term active, a batched cube
+    trace over every one of them -- 203 seconds for a single descent. The
+    pairwise term does not need any of that. Writing S for the alignment sum
+    and C for `permuted_text @ visual`, swapping positions p and q changes S
+    by 2(C_pq + C_qp - C_pp - C_qq + 2 A_pq B_pq), so one matrix product
+    yields the entire table.
+
+    Both fields must be symmetric with zero diagonal, which standardisation
+    already guarantees. Returns a dense [N, N] table of energy deltas; the
+    sign convention matches `ClosedFormEnergy.energy`, so negative is an
+    improvement.
+    """
+    size = permuted_text.shape[-1]
+    product = permuted_text @ visual
+    diagonal = torch.diagonal(product)
+    gain = 2.0 * (
+        product + product.T
+        - diagonal[:, None] - diagonal[None, :]
+        + 2.0 * permuted_text * visual
+    )
+    return -2.0 * gain / (size * (size - 1))
+
+
+def fast_pair_descent(
+    text: torch.Tensor, visual: torch.Tensor, start: torch.Tensor, max_steps: int
+) -> torch.Tensor:
+    """Steepest descent on the pairwise energy using the closed-form table."""
+    size = len(visual)
+    current = start.clone()
+    triangle = torch.triu(
+        torch.ones(size, size, dtype=torch.bool, device=visual.device), diagonal=1
+    )
+    for _ in range(max_steps):
+        permuted = text[current[:, None], current[None, :]]
+        deltas = all_pair_swap_deltas(permuted, visual)
+        deltas = torch.where(triangle, deltas, torch.inf)
+        best = int(deltas.argmin())
+        if float(deltas.view(-1)[best]) >= -1e-12:
+            break
+        p, q = best // size, best % size
+        current[[p, q]] = current[[q, p]]
+    return current