belka_gate_fusions: inject GateFusionResidual variants to asym and dual encoder

src/models/asymmetric_encoder.py

@@ -260,6 +260,9 @@ class TextFusionMLP(nn.Module):
 # Main model: AsymmetricEncoder
 # ---------------------------------------------------------------------------
 
+# ResidualGateFusin experiment
+from residual_fusions import ResidualGateType, GatedFusionResidual
+
 class AsymmetricEncoder(nn.Module):
     """Dual encoder for CVGL with text fusion on both branches.
 
@@ -294,6 +297,7 @@ class AsymmetricEncoder(nn.Module):
 
     def __init__(
         self,
+        gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate,
         dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
         dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
         lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
@@ -365,8 +369,12 @@ class AsymmetricEncoder(nn.Module):
             )
 
         # Separate gated fusion for query and gallery branches.
-        self.fusion_query = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
+        #! Experimental Gated fusion on query branch.
-        self.fusion_gallery = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
+        self.fusion_query = GatedFusionResidual(gate_type=gate_type,
+            init_gate=init_gate, baseline_mode=baseline_mode)
+        self.fusion_gallery = GatedFusionResidual(gate_type=gate_type,
+            init_gate=init_gate, baseline_mode=baseline_mode)
+
 
     @staticmethod
     def _freeze(module: nn.Module) -> None:
@@ -420,7 +428,7 @@ class AsymmetricEncoder(nn.Module):
         l1_texts: list[str] | None,
         l2_texts: list[str] | None,
         l3_texts: list[str] | None,
-        fusion: GatedFusion,
+        fusion: GatedFusionResidual,
     ) -> torch.Tensor:
         """Fuse image features with optional text, respecting per-sample presence.
 
@@ -451,8 +459,7 @@ class AsymmetricEncoder(nn.Module):
 
         # Per-sample fusion: text-present samples use full gated fusion,
         # empty-caption samples pass through pure image features.
-        gate = torch.sigmoid(fusion.alpha)
+        fused_with_text = fusion(img_feat, z_text)
-        fused_with_text = gate * img_feat + (1.0 - gate) * z_text
         out = torch.where(has_text.unsqueeze(-1), fused_with_text, img_feat)
         return F.normalize(out, dim=-1)
 

src/models/dual_encoder.py

@@ -21,6 +21,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+# residual fusions exp
+from residual_fusions import ResidualGateType, GatedFusionResidual
+
 
 class ProjectionHead(nn.Module):
     """MLP projection head with L2 normalization."""
@@ -83,120 +86,6 @@ 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.
@@ -222,6 +111,7 @@ class DualEncoderCaptionTest(nn.Module):
         baseline_mode: bool = False,
         init_gate: float = 0.7,
         device: str = "cuda",
+        gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate
     ) -> None:
         super().__init__()
         self.embed_dim = embed_dim
@@ -247,7 +137,11 @@ class DualEncoderCaptionTest(nn.Module):
             self._apply_unfreeze(unfreeze_mode)
 
         # Gated fusion on query branch.
-        self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
+        # self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
+
+        #! Experimental Gated fusion on query branch.
+        self.fusion = GatedFusionResidual(gate_type=gate_type,
+            init_gate=init_gate, baseline_mode=baseline_mode)
 
         # Projection heads.
         self.proj_query = ProjectionHead(

src/models/residual_fusions.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+import gin
+
+#! GATE-FUSIONS MODIFICATIONS ---------------------------------
+#! in_dim = 1024
+
+from enum import Enum
+import math
+
+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
+
+def init_bias_for_sigmoid(linear: nn.Linear, value: float) -> None:
+    nn.init.zeros_(linear.weight)
+    nn.init.constant_(linear.bias, math.log(value / (1.0 - value)))
+
+def init_residual_projs(linear: nn.Linear, scale: float) -> None:
+    nn.init.xavier_uniform_(linear.weight, gain=scale)
+    nn.init.zeros_(linear.bias)
+
+RESIDUAL_GATES = {
+    ResidualGateType.alpha_res_sum,
+    ResidualGateType.alpha_res_cat
+}
+
+@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 sum (per sample)
+    V5 - Alpha-weighted residual concat (per sample)
+    """
+
+    def __init__(self, gate_type: ResidualGateType, 
+                 init_gate: float = 0.7, in_dim = 1024, 
+                 init_res_weight: float = 0.1, residual_proj_scale: float = 0.1,  
+                 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
+        if gate_type == ResidualGateType.cross_gate:
+            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 
+        if gate_type in RESIDUAL_GATES:
+            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)
+            init_bias_for_sigmoid(self.weight_net_for_sum, value=init_res_weight)
+            init_bias_for_sigmoid(self.weight_net_for_cat, value=init_res_weight)
+            init_residual_projs(self.final_cat_residual_proj, scale=residual_proj_scale)
+
+        self.gate_type = gate_type
+        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
+
+        if self.gate_type == ResidualGateType.simple_residual_one_gate:
+            fused_vec = self.FuseSRGF(img_feat=img_feat, text_feat=text_feat)
+        if self.gate_type == ResidualGateType.cross_gate:
+            fused_vec = self.FuseRCGF(img_feat=img_feat, text_feat=text_feat)
+        if self.gate_type == ResidualGateType.gate_sum:
+            fused_vec = self.FuseGSUM(img_feat=img_feat, text_feat=text_feat)
+        if self.gate_type == ResidualGateType.alpha_res_sum:
+            fused_vec = self.FuseARGFSum(img_feat=img_feat, text_feat=text_feat)
+        if self.gate_type == ResidualGateType.alpha_res_cat:
+            fused_vec = self.FuseARGFCat(img_feat=img_feat, text_feat=text_feat)
+
+        return fused_vec

src/training/train_gtauav.py

@@ -99,7 +99,8 @@ class TrainConfigGTAUAV:
 
     # Training.
     resume_from: str | None = None  # path to checkpoint for resuming
-    output_dir: str = "out/gtauav/with_text"
+    # output_dir: str = "out/gtauav/with_text"
+    output_dir: str = "out/gtauav/with_text_exp_gate_SRGF"
     epochs: int = 10
     batch_size: int = 8
     num_workers: int = 4