From e2fd16e66c953cc707f3dded6ca9377799ab4d62 Mon Sep 17 00:00:00 2001
From: pikaliov <pikaliov@gmail.com>
Date: Fri, 15 May 2026 16:30:11 +0300
Subject: [PATCH] belka_gate_fusions: add gate-variants for CVGL exp.

---
 src/models/dual_encoder.py | 115 +++++++++++++++++++++++++++++++++++++
 1 file changed, 115 insertions(+)

diff --git a/src/models/dual_encoder.py b/src/models/dual_encoder.py
index e4e0b6b..558d318 100644
--- a/src/models/dual_encoder.py
+++ b/src/models/dual_encoder.py
@@ -46,6 +46,7 @@ class ProjectionHead(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return F.normalize(self.proj(x), dim=-1)
 
+#! GATE-FUSION ORIG --------------------------------- 
 
 @gin.configurable
 class GatedFusion(nn.Module):
@@ -82,6 +83,120 @@ class GatedFusion(nn.Module):
         return torch.sigmoid(self.alpha).item()
 
 
+#! GATE-FUSIONS MODIFICATIONS ---------------------------------
+#! in_dim = 1024
+
+from enum import Enum
+
+
+class ResidualGateType(Enum):
+    simple_residual_one_gate = 0,
+    cross_gate = 1,
+    gate_sum = 2,
+    alpha_res_cat = 3,
+    alpha_res_sum = 4
+
+# TODO: add GatedFusionresidual class to gin 
+@gin.configurable
+class GatedFusionresidual(nn.Module):
+    """Learnable gated fusion of image and text embeddings.
+
+    V1 - Simple residual gating with 1 common gate: 
+    V2 - Cross residual gating with 2 cross-gates
+    V3 - Gate + Simple Sum of feats x & y
+    V4 - Alpha-weighted residual concat (per sample)
+    V5 - Alpha-weighted residual sum (per sample)
+
+    """
+
+    def __init__(self, init_gate: float = 0.7, in_dim = 1024, baseline_mode: bool = False) -> None:
+        super().__init__()
+        # alpha is in logit space: sigmoid(alpha) = init_gate
+        init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
+        self.alpha = nn.Parameter(init_alpha)
+        
+        # alphas for separated cases
+        init_alpha_img_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
+        init_alpha_text_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
+        self.alpha_img = nn.Parameter(init_alpha_img_cross_gate)
+        self.alpha_text = nn.Parameter(init_alpha_text_cross_gate)
+        
+        # weight for sum and cat residual 
+        self.final_cat_residual_proj = nn.Linear(in_dim * 2, in_dim)
+        self.weight_net_for_sum = nn.Linear(in_dim, 1)
+        self.weight_net_for_cat = nn.Linear(in_dim * 2, 1)
+
+        self.baseline_mode = baseline_mode
+
+    def FuseSRGF(self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate = torch.sigmoid(self.alpha)
+        img_res = img_feat * gate + img_feat
+        text_res = text_feat * (1 - gate) + text_feat
+        fused_vec = img_res + text_res
+        return fused_vec
+    
+    def FuseRCGF(self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate_img = torch.sigmoid(self.alpha_img)
+        gate_text = torch.sigmoid(self.alpha_text)
+        z_img = img_feat + gate_text * img_feat
+        z_text = text_feat + gate_img * text_feat
+        fused_vec = z_img + z_text
+        return fused_vec
+    
+    def FuseGSUM(self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate = torch.sigmoid(self.alpha)
+        fuzed_vec = img_feat + text_feat + gate * img_feat + (1.0 - gate) * text_feat
+        return fuzed_vec
+    
+    def FuseARGFSum(
+        self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate = torch.sigmoid(self.alpha)
+        residual = img_feat + text_feat
+        res_weight = torch.sigmoid(self.weight_net_for_sum(residual))
+        fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
+        return fuzed_vec
+
+    def FuseARGFCat(
+        self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate = torch.sigmoid(self.alpha)
+        cat_vec = torch.cat([img_feat, text_feat], dim=-1)
+        residual = self.final_cat_residual_proj(cat_vec)
+        res_weight = torch.sigmoid(self.weight_net_for_cat(cat_vec))
+        fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
+        return fuzed_vec
+
+    def forward(
+        self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if text_feat is None or self.baseline_mode:
+            return img_feat
+        gate = torch.sigmoid(self.alpha)
+
+        # TODO: switch forwards here
+
+        return gate * img_feat + (1.0 - gate) * text_feat
+
+
+
+
+
 @gin.configurable
 class DualEncoderCaptionTest(nn.Module):
     """GeoRSCLIP dual encoder with gated text fusion on query branch.