Add tropical, compression, curvature, renormalization heads

heads/__init__.py

@@ -9,6 +9,9 @@ from .patch_attention.head import PatchAttention
 from .graph_crf.head import GraphCRF
 from .hypercolumn_linear.head import HypercolumnLinear
 from .info_bottleneck.head import InfoBottleneck
+from .tropical.head import TropicalSegmentation
+from .compression.head import CompressionSegmentation
+from .curvature.head import CurvatureSegmentation
 
 REGISTRY = {
     "linear_probe": LinearProbe,
@@ -20,6 +23,9 @@ REGISTRY = {
     "graph_crf": GraphCRF,
     "hypercolumn_linear": HypercolumnLinear,
     "info_bottleneck": InfoBottleneck,
+    "tropical": TropicalSegmentation,
+    "compression": CompressionSegmentation,
+    "curvature": CurvatureSegmentation,
 }
 
 ALL_NAMES = list(REGISTRY.keys())

heads/compression/__init__.py

@@ -0,0 +1 @@
+from .head import *

heads/compression/head.py

@@ -0,0 +1,26 @@
+"""Compression Segmentation: modulate features by local prediction residual.
+Patches that can't be predicted from neighbors get amplified before classification."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CompressionSegmentation(nn.Module):
+    name = "compression"
+    needs_intermediates = False
+
+    def __init__(self, feat_dim=768, num_classes=150):
+        super().__init__()
+        self.cls = nn.Conv2d(feat_dim, num_classes, 1)
+
+    def forward(self, spatial, inter=None):
+        B, C, H, W = spatial.shape
+        kernel = torch.tensor([[1, 1, 1], [1, 0, 1], [1, 1, 1]],
+                              dtype=spatial.dtype, device=spatial.device) / 8
+        kernel = kernel.reshape(1, 1, 3, 3).expand(C, 1, 3, 3)
+        neighbor_mean = F.conv2d(spatial, kernel, padding=1, groups=C)
+        surprise = (spatial - neighbor_mean).pow(2).sum(dim=1, keepdim=True)
+        surprise_norm = surprise / surprise.amax(dim=(2, 3), keepdim=True).clamp(min=1e-6)
+        modulated = spatial * (1 + surprise_norm * 3)
+        return self.cls(modulated)

heads/curvature/__init__.py

@@ -0,0 +1 @@
+from .head import *

heads/curvature/head.py

@@ -0,0 +1,38 @@
+"""Curvature Segmentation: modulate features by discrete Riemannian curvature.
+High-curvature locations (where the feature manifold bends sharply) are
+segment boundaries. Features at those locations get amplified."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CurvatureSegmentation(nn.Module):
+    name = "curvature"
+    needs_intermediates = False
+
+    def __init__(self, feat_dim=768, num_classes=150):
+        super().__init__()
+        self.cls = nn.Conv2d(feat_dim, num_classes, 1)
+
+    def _curvature(self, spatial):
+        B, C, H, W = spatial.shape
+        dx = torch.zeros_like(spatial)
+        dx[:, :, :, 1:-1] = (spatial[:, :, :, 2:] - spatial[:, :, :, :-2]) / 2
+        dy = torch.zeros_like(spatial)
+        dy[:, :, 1:-1, :] = (spatial[:, :, 2:, :] - spatial[:, :, :-2, :]) / 2
+        g11 = (dx * dx).sum(dim=1, keepdim=True)
+        g22 = (dy * dy).sum(dim=1, keepdim=True)
+        det_g = (g11 * g22).clamp(min=1e-10)
+        dxx = torch.zeros_like(spatial)
+        dxx[:, :, :, 1:-1] = spatial[:, :, :, 2:] - 2 * spatial[:, :, :, 1:-1] + spatial[:, :, :, :-2]
+        dyy = torch.zeros_like(spatial)
+        dyy[:, :, 1:-1, :] = spatial[:, :, 2:, :] - 2 * spatial[:, :, 1:-1, :] + spatial[:, :, :-2, :]
+        laplacian = dxx + dyy
+        return laplacian.pow(2).sum(dim=1, keepdim=True) / det_g
+
+    def forward(self, spatial, inter=None):
+        curv = self._curvature(spatial)
+        curv_norm = curv / curv.amax(dim=(2, 3), keepdim=True).clamp(min=1e-6)
+        modulated = spatial * (1 + curv_norm * 3)
+        return self.cls(modulated)

heads/tropical/__init__.py

@@ -0,0 +1 @@
+from .head import *

heads/tropical/head.py

@@ -0,0 +1,24 @@
+"""Tropical Segmentation: tropical inner product replaces standard dot product.
+logit_c = min_d(w_{c,d} + x_d) instead of sum(w_{c,d} * x_d)."""
+
+import torch
+import torch.nn as nn
+
+
+class TropicalSegmentation(nn.Module):
+    name = "tropical"
+    needs_intermediates = False
+
+    def __init__(self, feat_dim=768, num_classes=150):
+        super().__init__()
+        self.tropical_w = nn.Parameter(torch.randn(num_classes, feat_dim) * 0.01)
+        self.bias = nn.Parameter(torch.zeros(num_classes))
+
+    def forward(self, spatial, inter=None):
+        B, C, H, W = spatial.shape
+        f = spatial.permute(0, 2, 3, 1).reshape(-1, C)
+        # Tropical: min_d(w_{c,d} + f_d), smooth approx via -logsumexp(-beta*(w+f))/beta
+        beta = 10.0
+        expanded = self.tropical_w.unsqueeze(0) + f.unsqueeze(1)
+        logits = -torch.logsumexp(-beta * expanded, dim=2) / beta + self.bias
+        return logits.reshape(B, H, W, -1).permute(0, 3, 1, 2)