rpotter6298
78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
"""fusion_bridge — FusionBridge: N-input element-wise 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 element-wise.
|
|
|
|
Pure embedding producer — no classification head. Attach a head stage in
|
|
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
|
|
----------
|
|
input_dims : ordered list of input embedding dims
|
|
fusion_dim : projection / output dimension
|
|
additive : if True, use the shifted-multiply form (default: False)
|
|
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,
|
|
additive: bool = False,
|
|
use_se: bool = True,
|
|
se_reduction: int = 16,
|
|
se_pre_norm: bool = True,
|
|
):
|
|
super().__init__()
|
|
self.out_dim = fusion_dim
|
|
self.additive = additive
|
|
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)}"
|
|
)
|
|
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]))
|
|
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)
|