v4 update

v4/classes/bridges/fusion_bridge.py

@@ -0,0 +1,57 @@
+"""fusion_bridge — FusionBridge: N-input Hadamard-product fusion, embedding output only."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from v4.classes.accessory.se_block import SEBlock
+
+
+class FusionBridge(nn.Module):
+    """Project N input embeddings to a shared dim, fuse via element-wise product.
+
+    Pure embedding producer — no classification head.  Attach a head stage in
+    the pipeline config to produce logits.
+
+    Parameters
+    ----------
+    input_dims   : ordered list of input embedding dims
+    fusion_dim   : projection / output dimension
+    use_se       : SE gate on the fused vector
+    se_reduction : SE reduction factor
+    se_pre_norm  : LayerNorm before each projection; else Identity
+    """
+
+    def __init__(
+        self,
+        input_dims:   list[int],
+        fusion_dim:   int  = 256,
+        use_se:       bool = True,
+        se_reduction: int  = 16,
+        se_pre_norm:  bool = True,
+    ):
+        super().__init__()
+        self.out_dim = fusion_dim
+        self.W  = nn.ModuleList([nn.Linear(d, fusion_dim) for d in input_dims])
+        self.ln = nn.ModuleList(
+            [nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
+             for _ in input_dims]
+        )
+        self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None
+
+    def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
+        assert len(embeddings) == len(self.W), (
+            f"FusionBridge expects {len(self.W)} inputs, got {len(embeddings)}"
+        )
+        h = self.ln[0](self.W[0](embeddings[0]))
+        for i in range(1, len(embeddings)):
+            h = h * self.ln[i](self.W[i](embeddings[i]))
+        if self.se is not None:
+            h, _ = self.se(h)
+        return h
+
+    def set_phase(self, phase: str) -> None:
+        """Freeze bridge during tower_warmup; trainable otherwise."""
+        enabled = phase not in ("tower_warmup", "cd_warmup")
+        for p in self.parameters():
+            p.requires_grad_(enabled)

v4/classes/bridges/hyperbridge.py

@@ -0,0 +1,56 @@
+"""hyperbridge — HyperBridge: bilateral fusion over paired embeddings, embedding output only."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class HyperBridge(nn.Module):
+    """Fuse side embeddings (e.g. two z_fused vectors) into a single embedding.
+
+    Pure embedding producer — no classification head.  Attach a head stage in
+    the pipeline config to produce logits.
+
+    Modes
+    -----
+    embedding_mlp (default)
+        Linear projection of concatenated inputs → hidden_dim embedding.
+
+    classic_bridge
+        Per-side projection → Hadamard product → hidden_dim embedding.
+
+    Parameters
+    ----------
+    input_dims  : {side_key: dim} — e.g. {"a": 256, "b": 256}
+    hidden_dim  : output embedding dimension
+    mode        : "embedding_mlp" | "classic_bridge"
+    """
+
+    def __init__(
+        self,
+        input_dims: dict[str, int],
+        hidden_dim: int = 256,
+        mode:       str = "embedding_mlp",
+    ):
+        super().__init__()
+        self.input_names = list(input_dims.keys())
+        self.mode        = mode
+        self.out_dim     = hidden_dim
+        dims             = list(input_dims.values())
+
+        if mode == "embedding_mlp":
+            self.proj = nn.Linear(sum(dims), hidden_dim)
+        elif mode == "classic_bridge":
+            self.W  = nn.ModuleList([nn.Linear(d, hidden_dim) for d in dims])
+            self.ln = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in dims])
+        else:
+            raise ValueError(f"Unknown HyperBridge mode: {mode!r}")
+
+    def forward(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor:
+        ordered = [inputs[name] for name in self.input_names]
+        if self.mode == "embedding_mlp":
+            return self.proj(torch.cat(ordered, dim=1))
+        h = self.ln[0](self.W[0](ordered[0]))
+        for i in range(1, len(ordered)):
+            h = h * self.ln[i](self.W[i](ordered[i]))
+        return h