belka_gate_fusions: TODO: fix Incorrect fusion behaviour

conf/gtauav_balanced.gin

@@ -43,7 +43,7 @@ TrainConfigGTAUAV.dss_reembed_every = 1
 TrainConfigGTAUAV.use_mutex_sampler = True  # legacy flag, kept True unless disabling both samplers
 
 # ---- Output ----
-TrainConfigGTAUAV.output_dir = "out/gtauav/with_text"
+TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_exp_gate_SRGF"
 
 # ---- Tracking ----
 TrainConfigGTAUAV.use_wandb = False

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,8 @@ 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 +370,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)
-        self.fusion_gallery = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode)
+        #! Experimental Gated fusion on query branch.
+        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 +429,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 +460,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 = gate * img_feat + (1.0 - gate) * z_text
+        fused_with_text = fusion(img_feat, 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."""
@@ -46,6 +49,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):
@@ -107,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
@@ -132,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,154 @@
+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
+    
+    @property
+    def gate_value(self) -> float:
+        """Current gate value (image weight). 1.0 = text ignored."""
+        if self.baseline_mode:
+            return 1.0
+        return torch.sigmoid(self.alpha).item()
+    
+
+    

src/training/train.py

@@ -56,8 +56,8 @@ class TrainConfig:
         self,
         train_query_file: str = "Index/train_query.txt",
         val_query_file: str = "Index/test_query.txt",
-        data_root: str = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc",
-        output_dir: str = "out/caption_test",
+        data_root: str = "/media/servml/SSD_2_TB/datasets/cvgl_datasets/UAV-GeoLoc",
+        output_dir: str = "out/caption_test_exp_gate_SRGF",
         epochs: int = 10,
         batch_size: int = 128,
         num_workers: int = 4,

src/training/train_gtauav.py

@@ -58,8 +58,8 @@ from src.models.asymmetric_encoder import (
 LOGGER = logging.getLogger("caption_test.train_gtauav")
 
 # Default paths.
-_RGB_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR"
-_CAPTION_ROOT = "/home/servml/Документы/datasets/GTA-UAV-LR-captions"
+_RGB_ROOT = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR/"
+_CAPTION_ROOT = "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions/"
 _TRAIN_JSON = "meta/train_80.json"
 _TEST_JSON = "meta/test_20.json"
 
@@ -99,7 +99,9 @@ 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"
+    # ! ----------------- experimental gate outs path -----------------------
+    output_dir: str = "out/gtauav/with_text_exp_gate_SRGF"
     epochs: int = 10
     batch_size: int = 8
     num_workers: int = 4