update with orthobridge
This commit is contained in:
@@ -0,0 +1,55 @@
|
|||||||
|
"""concat_bridge — ConcatBridge: N-input concatenate + linear fusion.
|
||||||
|
|
||||||
|
The simplest possible fusion: glue all input embeddings end-to-end and let a
|
||||||
|
single Linear layer learn the mixing. No multiplicative interactions, no
|
||||||
|
zero-collapse risk, no cross-term explosion as N grows.
|
||||||
|
|
||||||
|
Useful as a baseline against the multiplicative bridges (FusionBridge,
|
||||||
|
PairwiseAdditiveBridge): if this gets within noise of them, then multiplicative
|
||||||
|
fusion isn't actually buying us anything.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
from v4.classes.accessory.se_block import SEBlock
|
||||||
|
|
||||||
|
|
||||||
|
class ConcatBridge(nn.Module):
|
||||||
|
"""Concatenate N input embeddings, project down to fusion_dim.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
input_dims : ordered list of input embedding dims
|
||||||
|
fusion_dim : output dimension after the linear projection
|
||||||
|
use_ln : LayerNorm after the projection (default: True)
|
||||||
|
use_se : SE gate on the fused vector
|
||||||
|
se_reduction : SE bottleneck factor
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
input_dims: list[int],
|
||||||
|
fusion_dim: int = 256,
|
||||||
|
use_ln: bool = True,
|
||||||
|
use_se: bool = True,
|
||||||
|
se_reduction: int = 16,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.out_dim = fusion_dim
|
||||||
|
self.proj = nn.Linear(sum(input_dims), fusion_dim)
|
||||||
|
self.ln = nn.LayerNorm(fusion_dim) if use_ln else nn.Identity()
|
||||||
|
self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None
|
||||||
|
|
||||||
|
def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
|
||||||
|
h = torch.cat(embeddings, dim=-1)
|
||||||
|
h = self.ln(self.proj(h))
|
||||||
|
if self.se is not None:
|
||||||
|
h, _ = self.se(h)
|
||||||
|
return h
|
||||||
|
|
||||||
|
def set_phase(self, phase: str) -> None:
|
||||||
|
enabled = phase not in ("tower_warmup", "cd_warmup")
|
||||||
|
for p in self.parameters():
|
||||||
|
p.requires_grad_(enabled)
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
"""fusion_bridge — FusionBridge: N-input Hadamard-product fusion, embedding output only."""
|
"""fusion_bridge — FusionBridge: N-input element-wise fusion, embedding output only."""
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
@@ -8,15 +8,27 @@ from v4.classes.accessory.se_block import SEBlock
|
|||||||
|
|
||||||
|
|
||||||
class FusionBridge(nn.Module):
|
class FusionBridge(nn.Module):
|
||||||
"""Project N input embeddings to a shared dim, fuse via element-wise product.
|
"""Project N input embeddings to a shared dim, fuse element-wise.
|
||||||
|
|
||||||
Pure embedding producer — no classification head. Attach a head stage in
|
Pure embedding producer — no classification head. Attach a head stage in
|
||||||
the pipeline config to produce logits.
|
the pipeline config to produce logits.
|
||||||
|
|
||||||
|
Two fusion modes:
|
||||||
|
* additive=False (Hadamard): h = ∏_i ln_i(W_i z_i)
|
||||||
|
Interaction term only. A near-zero factor
|
||||||
|
in any stream silences that dimension for
|
||||||
|
the whole bridge.
|
||||||
|
* additive=True (shifted) : h = ∏_i (1 + ln_i(W_i z_i)) - 1
|
||||||
|
Expands to Σ_i a_i + cross-terms (sums of
|
||||||
|
products). A silent stream (≈0) reduces to
|
||||||
|
identity on its factor, so other streams'
|
||||||
|
contributions survive unchanged.
|
||||||
|
|
||||||
Parameters
|
Parameters
|
||||||
----------
|
----------
|
||||||
input_dims : ordered list of input embedding dims
|
input_dims : ordered list of input embedding dims
|
||||||
fusion_dim : projection / output dimension
|
fusion_dim : projection / output dimension
|
||||||
|
additive : if True, use the shifted-multiply form (default: False)
|
||||||
use_se : SE gate on the fused vector
|
use_se : SE gate on the fused vector
|
||||||
se_reduction : SE reduction factor
|
se_reduction : SE reduction factor
|
||||||
se_pre_norm : LayerNorm before each projection; else Identity
|
se_pre_norm : LayerNorm before each projection; else Identity
|
||||||
@@ -26,12 +38,14 @@ class FusionBridge(nn.Module):
|
|||||||
self,
|
self,
|
||||||
input_dims: list[int],
|
input_dims: list[int],
|
||||||
fusion_dim: int = 256,
|
fusion_dim: int = 256,
|
||||||
|
additive: bool = False,
|
||||||
use_se: bool = True,
|
use_se: bool = True,
|
||||||
se_reduction: int = 16,
|
se_reduction: int = 16,
|
||||||
se_pre_norm: bool = True,
|
se_pre_norm: bool = True,
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.out_dim = fusion_dim
|
self.out_dim = fusion_dim
|
||||||
|
self.additive = additive
|
||||||
self.W = nn.ModuleList([nn.Linear(d, fusion_dim) for d in input_dims])
|
self.W = nn.ModuleList([nn.Linear(d, fusion_dim) for d in input_dims])
|
||||||
self.ln = nn.ModuleList(
|
self.ln = nn.ModuleList(
|
||||||
[nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
|
[nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
|
||||||
@@ -43,6 +57,12 @@ class FusionBridge(nn.Module):
|
|||||||
assert len(embeddings) == len(self.W), (
|
assert len(embeddings) == len(self.W), (
|
||||||
f"FusionBridge expects {len(self.W)} inputs, got {len(embeddings)}"
|
f"FusionBridge expects {len(self.W)} inputs, got {len(embeddings)}"
|
||||||
)
|
)
|
||||||
|
if self.additive:
|
||||||
|
h = 1.0 + self.ln[0](self.W[0](embeddings[0]))
|
||||||
|
for i in range(1, len(embeddings)):
|
||||||
|
h = h * (1.0 + self.ln[i](self.W[i](embeddings[i])))
|
||||||
|
h = h - 1.0
|
||||||
|
else:
|
||||||
h = self.ln[0](self.W[0](embeddings[0]))
|
h = self.ln[0](self.W[0](embeddings[0]))
|
||||||
for i in range(1, len(embeddings)):
|
for i in range(1, len(embeddings)):
|
||||||
h = h * self.ln[i](self.W[i](embeddings[i]))
|
h = h * self.ln[i](self.W[i](embeddings[i]))
|
||||||
|
|||||||
@@ -0,0 +1,80 @@
|
|||||||
|
"""gated_bridge — GatedAdditiveBridge.
|
||||||
|
|
||||||
|
Per-sample, per-stream learned gates determine how much each stream contributes
|
||||||
|
to the fused embedding. Each gate is a sigmoid scalar produced by an MLP over
|
||||||
|
the raw input embeddings, so the gate is conditioned on the actual content of
|
||||||
|
all streams — when a stream's signal is weak for a particular sample, its gate
|
||||||
|
can attenuate toward 0; when it's informative, gate goes toward 1.
|
||||||
|
|
||||||
|
h = Σ_i g_i(x) · LN_i(W_i z_i)
|
||||||
|
g_i(x) = σ(MLP_i([z_1, z_2, …, z_N]))
|
||||||
|
|
||||||
|
No symmetry-breaking between streams — every tower is treated identically; the
|
||||||
|
gate network decides per-sample which to amplify. Gates use sigmoid (not
|
||||||
|
softmax) so they can be independently small or large; the model isn't forced
|
||||||
|
into a "pick one" distribution.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
from v4.classes.accessory.se_block import SEBlock
|
||||||
|
|
||||||
|
|
||||||
|
class GatedAdditiveBridge(nn.Module):
|
||||||
|
"""Per-sample sigmoid-gated additive fusion.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
input_dims : ordered list of input embedding dims
|
||||||
|
fusion_dim : projection / output dimension
|
||||||
|
gate_hidden : hidden width of the gating MLP (default: fusion_dim)
|
||||||
|
use_ln : LayerNorm after each per-stream projection (default: True)
|
||||||
|
use_se : SE gate on the fused vector
|
||||||
|
se_reduction : SE bottleneck factor
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
input_dims: list[int],
|
||||||
|
fusion_dim: int = 256,
|
||||||
|
gate_hidden: int = 128,
|
||||||
|
use_ln: bool = True,
|
||||||
|
use_se: bool = True,
|
||||||
|
se_reduction: int = 16,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.out_dim = fusion_dim
|
||||||
|
self.n = len(input_dims)
|
||||||
|
self.W = nn.ModuleList([nn.Linear(d, fusion_dim) for d in input_dims])
|
||||||
|
self.ln = nn.ModuleList(
|
||||||
|
[nn.LayerNorm(fusion_dim) if use_ln else nn.Identity()
|
||||||
|
for _ in input_dims]
|
||||||
|
)
|
||||||
|
self.gate = nn.Sequential(
|
||||||
|
nn.Linear(sum(input_dims), gate_hidden),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
nn.Linear(gate_hidden, self.n),
|
||||||
|
nn.Sigmoid(),
|
||||||
|
)
|
||||||
|
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) == self.n, (
|
||||||
|
f"GatedAdditiveBridge expects {self.n} inputs, got {len(embeddings)}"
|
||||||
|
)
|
||||||
|
projected = [self.ln[i](self.W[i](e)) for i, e in enumerate(embeddings)]
|
||||||
|
gate_in = torch.cat(embeddings, dim=-1)
|
||||||
|
gates = self.gate(gate_in) # (B, N) ∈ (0, 1)
|
||||||
|
h = 0
|
||||||
|
for i in range(self.n):
|
||||||
|
h = h + gates[..., i:i+1] * projected[i]
|
||||||
|
if self.se is not None:
|
||||||
|
h, _ = self.se(h)
|
||||||
|
return h
|
||||||
|
|
||||||
|
def set_phase(self, phase: str) -> None:
|
||||||
|
enabled = phase not in ("tower_warmup", "cd_warmup")
|
||||||
|
for p in self.parameters():
|
||||||
|
p.requires_grad_(enabled)
|
||||||
@@ -0,0 +1,141 @@
|
|||||||
|
"""ortho_bridge — OrthoBridge: wraps any underlying bridge with cross-tower
|
||||||
|
orthogonality regularization.
|
||||||
|
|
||||||
|
Forces each tower to encode information that the other towers DON'T encode by
|
||||||
|
penalising cross-tower representational similarity. The penalty is added to
|
||||||
|
the main training loss via the `modify_loss` hook, so the fusion stage runner
|
||||||
|
sees it transparently — no other code changes required.
|
||||||
|
|
||||||
|
How it works
|
||||||
|
------------
|
||||||
|
1. Inner bridge fuses the embeddings into the usual single tensor (forward).
|
||||||
|
2. While forward is running, we compute pairwise linear-CKA between every
|
||||||
|
pair of input embeddings and stash the average value as `self._stashed`.
|
||||||
|
3. `modify_loss(loss)` returns `loss + ortho_weight * stashed`. Gradients
|
||||||
|
from the penalty flow back through the embeddings into the tower weights,
|
||||||
|
pushing each tower's representations apart.
|
||||||
|
|
||||||
|
CKA reference: Kornblith et al., "Similarity of Neural Network Representations
|
||||||
|
Revisited" (ICML 2019). Linear CKA on centred embeddings is in [0, 1]:
|
||||||
|
0 = orthogonal (uncorrelated) representations
|
||||||
|
1 = identical (up to linear transform)
|
||||||
|
|
||||||
|
Config example:
|
||||||
|
{
|
||||||
|
"name": "nt",
|
||||||
|
"type": "fusion",
|
||||||
|
"module": "v4.classes.bridges.ortho_bridge",
|
||||||
|
"class": "OrthoBridge",
|
||||||
|
"args": {
|
||||||
|
"ortho_weight": 0.1,
|
||||||
|
"inner_module": "v4.classes.bridges.fusion_bridge",
|
||||||
|
"inner_class": "FusionBridge",
|
||||||
|
"inner_args": { "fusion_dim": 256, "use_se": true }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import importlib
|
||||||
|
from itertools import combinations
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
|
||||||
|
def _linear_cka(z_a: torch.Tensor, z_b: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
|
||||||
|
"""Linear CKA between two (B, D_a) and (B, D_b) embedding batches.
|
||||||
|
|
||||||
|
Centred Gram-matrix similarity, normalised to [0, 1].
|
||||||
|
"""
|
||||||
|
a = z_a - z_a.mean(dim=0, keepdim=True)
|
||||||
|
b = z_b - z_b.mean(dim=0, keepdim=True)
|
||||||
|
Ga = a @ a.T
|
||||||
|
Gb = b @ b.T
|
||||||
|
num = (Ga * Gb).sum()
|
||||||
|
den = torch.sqrt((Ga * Ga).sum() * (Gb * Gb).sum() + eps)
|
||||||
|
return num / den
|
||||||
|
|
||||||
|
|
||||||
|
class OrthoBridge(nn.Module):
|
||||||
|
"""Wrap any underlying bridge and add cross-tower orthogonality.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
input_dims : ordered list of input embedding dims (passed to inner bridge)
|
||||||
|
fusion_dim : convenience alias passed to inner bridge if it accepts it
|
||||||
|
ortho_weight : λ on the orthogonality penalty (default 0.1)
|
||||||
|
inner_module : import path of the inner bridge class
|
||||||
|
inner_class : class name within `inner_module`
|
||||||
|
inner_args : kwargs forwarded to the inner bridge constructor
|
||||||
|
|
||||||
|
The OrthoBridge does NOT touch the inner bridge's forward output — it only
|
||||||
|
computes and stashes a penalty during forward, then exposes it via the
|
||||||
|
`modify_loss` hook.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
input_dims: list[int],
|
||||||
|
fusion_dim: int = 256,
|
||||||
|
ortho_weight: float = 0.1,
|
||||||
|
inner_module: str = "v4.classes.bridges.fusion_bridge",
|
||||||
|
inner_class: str = "FusionBridge",
|
||||||
|
inner_args: dict | None = None,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.ortho_weight = float(ortho_weight)
|
||||||
|
|
||||||
|
# Build the inner bridge. Pass fusion_dim through unless the caller's
|
||||||
|
# inner_args overrides it.
|
||||||
|
inner_kwargs: dict[str, Any] = dict(inner_args or {})
|
||||||
|
inner_kwargs.setdefault("fusion_dim", fusion_dim)
|
||||||
|
mod = importlib.import_module(inner_module)
|
||||||
|
cls = getattr(mod, inner_class)
|
||||||
|
self.inner = cls(input_dims, **inner_kwargs)
|
||||||
|
self.out_dim = self.inner.out_dim
|
||||||
|
|
||||||
|
# Stash penalty here on every forward; modify_loss reads from it.
|
||||||
|
self.register_buffer("_stashed", torch.zeros(()), persistent=False)
|
||||||
|
self._last_pairs_cka: list[float] = []
|
||||||
|
|
||||||
|
# ── core ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
|
||||||
|
# Compute pairwise linear-CKA across the raw tower embeddings.
|
||||||
|
# We use the RAW per-tower embeddings (not the inner bridge's projected
|
||||||
|
# versions) because we want to push the TOWERS apart, not the bridge's
|
||||||
|
# internal projections.
|
||||||
|
if len(embeddings) >= 2:
|
||||||
|
ckas = []
|
||||||
|
for i, j in combinations(range(len(embeddings)), 2):
|
||||||
|
ckas.append(_linear_cka(embeddings[i], embeddings[j]))
|
||||||
|
penalty = torch.stack(ckas).mean()
|
||||||
|
self._stashed = penalty
|
||||||
|
self._last_pairs_cka = [float(c.detach().cpu()) for c in ckas]
|
||||||
|
else:
|
||||||
|
self._stashed = torch.zeros((), device=embeddings[0].device)
|
||||||
|
self._last_pairs_cka = []
|
||||||
|
|
||||||
|
return self.inner(embeddings)
|
||||||
|
|
||||||
|
def modify_loss(self, loss: torch.Tensor) -> torch.Tensor:
|
||||||
|
return loss + self.ortho_weight * self._stashed
|
||||||
|
|
||||||
|
# ── delegate to inner ───────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def set_phase(self, phase: str) -> None:
|
||||||
|
if hasattr(self.inner, "set_phase"):
|
||||||
|
self.inner.set_phase(phase)
|
||||||
|
|
||||||
|
# ── introspection ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
@property
|
||||||
|
def last_cka(self) -> float:
|
||||||
|
"""Mean cross-tower linear CKA from the most recent forward pass.
|
||||||
|
|
||||||
|
Useful for logging — should DECREASE during training if the penalty
|
||||||
|
is doing its job.
|
||||||
|
"""
|
||||||
|
return float(self._stashed.detach().cpu()) if self._stashed.numel() else 0.0
|
||||||
@@ -0,0 +1,77 @@
|
|||||||
|
"""pairwise_bridge — PairwiseAdditiveBridge.
|
||||||
|
|
||||||
|
For N input streams, compute the shifted-multiply ((1+a)(1+b) - 1) fusion for
|
||||||
|
every pair, then mix them with learnable per-pair scalar weights.
|
||||||
|
|
||||||
|
Motivation: the basic additive form (1+a)(1+b)...(1+N) - 1 helps for 2 streams
|
||||||
|
but regresses for 3+ because the triple-and-higher cross-terms (abc, abcd, …)
|
||||||
|
have explosive variance. This bridge keeps only the 2-way interactions —
|
||||||
|
O(N²) pair-experts rather than O(2^N) cross-terms — and lets the model learn
|
||||||
|
which pairs matter.
|
||||||
|
|
||||||
|
For 2 streams this reduces to a single weighted (1+a)(1+b)-1 → equivalent (up
|
||||||
|
to the scaling factor) to FusionBridge(additive=True).
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from itertools import combinations
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
from v4.classes.accessory.se_block import SEBlock
|
||||||
|
|
||||||
|
|
||||||
|
class PairwiseAdditiveBridge(nn.Module):
|
||||||
|
"""Sum of per-pair shifted-multiplies.
|
||||||
|
|
||||||
|
Parameters
|
||||||
|
----------
|
||||||
|
input_dims : ordered list of input embedding dims
|
||||||
|
fusion_dim : projection / output dimension
|
||||||
|
use_ln : LayerNorm after each per-stream projection (default: True)
|
||||||
|
use_se : SE gate on the fused vector
|
||||||
|
se_reduction : SE bottleneck factor
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
input_dims: list[int],
|
||||||
|
fusion_dim: int = 256,
|
||||||
|
use_ln: bool = True,
|
||||||
|
use_se: bool = True,
|
||||||
|
se_reduction: int = 16,
|
||||||
|
):
|
||||||
|
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 use_ln else nn.Identity()
|
||||||
|
for _ in input_dims]
|
||||||
|
)
|
||||||
|
n_pairs = max(1, len(input_dims) * (len(input_dims) - 1) // 2)
|
||||||
|
# initialise to uniform mixing so each pair contributes equally at start
|
||||||
|
self.pair_weights = nn.Parameter(torch.full((n_pairs,), 1.0 / n_pairs))
|
||||||
|
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"PairwiseAdditiveBridge expects {len(self.W)} inputs, got {len(embeddings)}"
|
||||||
|
)
|
||||||
|
projected = [self.ln[i](self.W[i](e)) for i, e in enumerate(embeddings)]
|
||||||
|
if len(projected) == 1:
|
||||||
|
h = projected[0]
|
||||||
|
else:
|
||||||
|
pairs = list(combinations(range(len(projected)), 2))
|
||||||
|
h = 0
|
||||||
|
for idx, (i, j) in enumerate(pairs):
|
||||||
|
pair_fusion = (1.0 + projected[i]) * (1.0 + projected[j]) - 1.0
|
||||||
|
h = h + self.pair_weights[idx] * pair_fusion
|
||||||
|
if self.se is not None:
|
||||||
|
h, _ = self.se(h)
|
||||||
|
return h
|
||||||
|
|
||||||
|
def set_phase(self, phase: str) -> None:
|
||||||
|
enabled = phase not in ("tower_warmup", "cd_warmup")
|
||||||
|
for p in self.parameters():
|
||||||
|
p.requires_grad_(enabled)
|
||||||
@@ -241,6 +241,8 @@ def run(
|
|||||||
if not losses:
|
if not losses:
|
||||||
continue
|
continue
|
||||||
loss = sum(losses) / len(losses)
|
loss = sum(losses) / len(losses)
|
||||||
|
if hasattr(bridge, "modify_loss"):
|
||||||
|
loss = bridge.modify_loss(loss)
|
||||||
opt.zero_grad(); loss.backward(); opt.step()
|
opt.zero_grad(); loss.backward(); opt.step()
|
||||||
total_loss += loss.item() * len(y_t)
|
total_loss += loss.item() * len(y_t)
|
||||||
total_n += len(y_t)
|
total_n += len(y_t)
|
||||||
@@ -254,6 +256,8 @@ def run(
|
|||||||
if logits is None:
|
if logits is None:
|
||||||
continue
|
continue
|
||||||
loss = F.cross_entropy(logits, y_t, weight=cw)
|
loss = F.cross_entropy(logits, y_t, weight=cw)
|
||||||
|
if hasattr(bridge, "modify_loss"):
|
||||||
|
loss = bridge.modify_loss(loss)
|
||||||
opt.zero_grad(); loss.backward(); opt.step()
|
opt.zero_grad(); loss.backward(); opt.step()
|
||||||
total_correct += int((logits.argmax(1) == y_t).sum())
|
total_correct += int((logits.argmax(1) == y_t).sum())
|
||||||
total_loss += loss.item() * len(y_t)
|
total_loss += loss.item() * len(y_t)
|
||||||
|
|||||||
@@ -309,6 +309,8 @@ def _parallel_fusion(
|
|||||||
if not losses:
|
if not losses:
|
||||||
continue
|
continue
|
||||||
loss = sum(losses) / len(losses)
|
loss = sum(losses) / len(losses)
|
||||||
|
if hasattr(bridge, "modify_loss"):
|
||||||
|
loss = bridge.modify_loss(loss)
|
||||||
ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
|
ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
|
||||||
total_loss += loss.item() * len(y_t)
|
total_loss += loss.item() * len(y_t)
|
||||||
total_n += len(y_t)
|
total_n += len(y_t)
|
||||||
@@ -322,6 +324,8 @@ def _parallel_fusion(
|
|||||||
if logits is None:
|
if logits is None:
|
||||||
continue
|
continue
|
||||||
loss = F.cross_entropy(logits, y_t, weight=ctx["class_weights"])
|
loss = F.cross_entropy(logits, y_t, weight=ctx["class_weights"])
|
||||||
|
if hasattr(bridge, "modify_loss"):
|
||||||
|
loss = bridge.modify_loss(loss)
|
||||||
ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
|
ctx["opt"].zero_grad(); loss.backward(); ctx["opt"].step()
|
||||||
total_correct += int((logits.argmax(1) == y_t).sum())
|
total_correct += int((logits.argmax(1) == y_t).sum())
|
||||||
total_loss += loss.item() * len(y_t)
|
total_loss += loss.item() * len(y_t)
|
||||||
|
|||||||
Binary file not shown.
@@ -0,0 +1,12 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "img+cd ensemble with additive nt bridge (shifted-multiply, no Hadamard collapse). 3 reps.",
|
||||||
|
"run_name": "experiments/tri_v1/additive/ensemble",
|
||||||
|
"reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"args": { "fusion_dim": 256, "additive": true }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "Promote additive ensemble to 10 reps. First 3 will be skipped (results exist).",
|
||||||
|
"run_name": "experiments/tri_v1/additive/ensemble",
|
||||||
|
"reps": 10,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"args": { "fusion_dim": 256, "additive": true }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "img+cd + UNet-CDR vector injection with additive nt bridge. Direct test of whether Hadamard suppression explained why geom inject doesn't help. 3 reps.",
|
||||||
|
"run_name": "experiments/tri_v1/additive/geom_vec_unet",
|
||||||
|
"reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"args": { "fusion_dim": 256, "additive": true }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "Full img+cd+geom tritower with additive nt bridge. With 3 streams competing through Hadamard, suppression risk is highest here — biggest expected lift if the hypothesis holds. 3 reps.",
|
||||||
|
"run_name": "experiments/tri_v1/additive/tritower",
|
||||||
|
"reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"args": { "fusion_dim": 256, "additive": true }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "Tritower with ConcatBridge — no multiplicative interactions baseline. 3 reps.",
|
||||||
|
"run_name": "experiments/tri_v1/altbridge/concat",
|
||||||
|
"reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.concat_bridge",
|
||||||
|
"class": "ConcatBridge",
|
||||||
|
"args": { "fusion_dim": 256 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "img+cd ensemble with ConcatBridge — 2-stream variant of the concat experiment. 3 reps.",
|
||||||
|
"run_name": "experiments/tri_v1/altbridge/ensemble_concat",
|
||||||
|
"reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.concat_bridge",
|
||||||
|
"class": "ConcatBridge",
|
||||||
|
"args": { "fusion_dim": 256 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "Tritower with GatedAdditiveBridge — per-sample sigmoid gates over each stream. 3 reps.",
|
||||||
|
"run_name": "experiments/tri_v1/altbridge/gated",
|
||||||
|
"reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.gated_bridge",
|
||||||
|
"class": "GatedAdditiveBridge",
|
||||||
|
"args": { "fusion_dim": 256 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "img+cd ensemble with GatedAdditiveBridge — per-sample sigmoid gates over each stream. 3 reps.",
|
||||||
|
"run_name": "experiments/tri_v1/altbridge/ensemble_gated",
|
||||||
|
"reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.gated_bridge",
|
||||||
|
"class": "GatedAdditiveBridge",
|
||||||
|
"args": { "fusion_dim": 256 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "Tritower with PairwiseAdditiveBridge — keeps 2-way interactions, drops 3-way/abc terms. 3 reps.",
|
||||||
|
"run_name": "experiments/tri_v1/altbridge/pairwise",
|
||||||
|
"reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.pairwise_bridge",
|
||||||
|
"class": "PairwiseAdditiveBridge",
|
||||||
|
"args": { "fusion_dim": 256 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"_note": "img+cd ensemble with PairwiseAdditiveBridge — for N=2 this reduces to a single weighted (1+a)(1+b)-1 fusion (≈ FusionBridge additive=True). 3 reps.",
|
||||||
|
"run_name": "experiments/tri_v1/altbridge/ensemble_pairwise",
|
||||||
|
"reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.pairwise_bridge",
|
||||||
|
"class": "PairwiseAdditiveBridge",
|
||||||
|
"args": { "fusion_dim": 256 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,36 @@
|
|||||||
|
[
|
||||||
|
{ "_note": "Bottleneck experiment for the 2-stream ensemble. 4 bridges × 3 reps at fusion_dim=8.",
|
||||||
|
|
||||||
|
"run_name": "experiments/tri_v1/bottleneck/hadamard_ensemble", "reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": { "args": { "fusion_dim": 8 } }
|
||||||
|
}
|
||||||
|
},
|
||||||
|
{ "run_name": "experiments/tri_v1/bottleneck/concat_ensemble", "reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.concat_bridge",
|
||||||
|
"class": "ConcatBridge",
|
||||||
|
"args": { "fusion_dim": 8 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
{ "run_name": "experiments/tri_v1/bottleneck/pairwise_ensemble", "reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.pairwise_bridge",
|
||||||
|
"class": "PairwiseAdditiveBridge",
|
||||||
|
"args": { "fusion_dim": 8 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
{ "run_name": "experiments/tri_v1/bottleneck/gated_ensemble", "reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.gated_bridge",
|
||||||
|
"class": "GatedAdditiveBridge",
|
||||||
|
"args": { "fusion_dim": 8 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,36 @@
|
|||||||
|
[
|
||||||
|
{ "_note": "Bottleneck experiment — fusion_dim=8 forces each bridge to compress aggressively. 4 bridges × 3 reps; if shapes-are-just-equivalent hypothesis is right, all 4 should still converge; if the bridges differ in compression priorities, AUC should diverge meaningfully.",
|
||||||
|
|
||||||
|
"run_name": "experiments/tri_v1/bottleneck/hadamard_tritower", "reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": { "args": { "fusion_dim": 8 } }
|
||||||
|
}
|
||||||
|
},
|
||||||
|
{ "run_name": "experiments/tri_v1/bottleneck/concat_tritower", "reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.concat_bridge",
|
||||||
|
"class": "ConcatBridge",
|
||||||
|
"args": { "fusion_dim": 8 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
{ "run_name": "experiments/tri_v1/bottleneck/pairwise_tritower", "reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.pairwise_bridge",
|
||||||
|
"class": "PairwiseAdditiveBridge",
|
||||||
|
"args": { "fusion_dim": 8 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
{ "run_name": "experiments/tri_v1/bottleneck/gated_tritower", "reps": 3,
|
||||||
|
"stage_overrides": {
|
||||||
|
"nt": {
|
||||||
|
"module": "v4.classes.bridges.gated_bridge",
|
||||||
|
"class": "GatedAdditiveBridge",
|
||||||
|
"args": { "fusion_dim": 8 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
Reference in New Issue
Block a user