Switch to shared DINOv3 WEB encoder (saves ~4-5 GB VRAM)

- Single DINOv3 WEB for both drone and satellite branches (shared_encoder=True default)
- One set of MONA adapters instead of two: 7M trainable vs 14M
- Total params: 438M (was 748M), trainable: 10.6M (was 17.6M)
- Asymmetric mode still available via shared_encoder=False
- Add gradient accumulation (grad_accum_steps, --grad-accum CLI flag)
- Update model summary in README

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

README.md

@@ -254,16 +254,18 @@ Mixed precision: AMP fp16 for model forward, fp32 for loss
 
 | Component | Params | Trainable | Notes |
 |-----------|--------|-----------|-------|
-| DINOv3 ViT-L/16 LVD (drone) | 303M | frozen | backbone weights frozen |
+| DINOv3 ViT-L/16 WEB (shared) | 303M | frozen | single encoder for drone + satellite |
-| MONA adapters (drone) | 7.0M | 7.0M | 2 per block × 24 blocks, bottleneck=64 |
+| MONA adapters (shared) | 7.0M | 7.0M | 2 per block × 24 blocks, bottleneck=64 |
-| DINOv3 ViT-L/16 SAT (satellite) | 303M | frozen | backbone weights frozen |
-| MONA adapters (satellite) | 7.0M | 7.0M | 2 per block × 24 blocks, bottleneck=64 |
 | DGTRS-CLIP ViT-L-14 (text) | 124M | frozen | backbone weights frozen |
 | LoRA adapters (text) | 147K | 147K | Q+V, rank=4, 12 blocks |
 | TextFusionMLP (shared) | 3.4M | 3.4M | Linear(2304,1024) + GELU + Linear(1024,1024) |
 | GatedFusion α_q + α_g | 2 | 2 | separate gate scalars |
 | logit_scale | 1 | 1 | learnable temperature |
-| **Total** | **748M** | **17.6M (2.35%)** | retrieval dim = 1024 |
+| **Total (shared)** | **438M** | **10.6M (2.42%)** | retrieval dim = 1024 |
+
+> **Asymmetric mode** (`--shared-encoder false`): uses separate DINOv3 WEB (drone) + DINOv3 SAT
+> (satellite) encoders with independent MONA adapters. Total: 748M params, 17.6M trainable.
+> Requires ~4-5 GB more VRAM.
 
 ## Experiments
 

conf/gtauav_balanced.gin

@@ -25,6 +25,7 @@ TrainConfigGTAUAV.device = "cuda"
 # ---- Model ----
 TrainConfigGTAUAV.init_gate = 0.7
 TrainConfigGTAUAV.baseline_mode = False
+TrainConfigGTAUAV.shared_encoder = True
 
 # ---- Loss ----
 TrainConfigGTAUAV.tau_init = 0.07

src/models/asymmetric_encoder.py

@@ -249,10 +249,16 @@ class TextFusionMLP(nn.Module):
 # ---------------------------------------------------------------------------
 
 class AsymmetricEncoder(nn.Module):
-    """Asymmetric dual encoder for CVGL with text fusion on both branches.
+    """Dual encoder for CVGL with text fusion on both branches.
 
-    Query branch:   DINOv3 LVD (drone) + text(L1/L2/L3) -> GatedFusion_q -> query [1024]
+    Supports two modes:
-    Gallery branch: DINOv3 SAT (sat)   + text(L1/L2/L3) -> GatedFusion_g -> gallery [1024]
+      - **shared** (default): single DINOv3 WEB encoder for both drone and satellite,
+        one set of MONA adapters. Saves ~4-5 GB VRAM and halves adapter params.
+      - **asymmetric**: separate DINOv3 encoders (LVD for drone, SAT for satellite),
+        each with their own MONA adapters (legacy mode).
+
+    Query branch:   DINOv3 (drone) + text(L1/L2/L3) -> GatedFusion_q -> query [1024]
+    Gallery branch: DINOv3 (sat)   + text(L1/L2/L3) -> GatedFusion_g -> gallery [1024]
 
     No projection layers — retrieval space is DINOv3 native 1024-dim.
     Text fusion MLP is shared between branches (same caption format).
@@ -262,11 +268,12 @@ class AsymmetricEncoder(nn.Module):
     (text_feat=None → gate acts as identity).
 
     Args:
-        dino_web_path: Path to DINOv3 LVD checkpoint (drone encoder).
+        dino_web_path: Path to DINOv3 LVD checkpoint (used for both branches in shared mode).
-        dino_sat_path: Path to DINOv3 SAT checkpoint (satellite encoder).
+        dino_sat_path: Path to DINOv3 SAT checkpoint (only used in asymmetric mode).
         lrsclip_path: Path to DGTRS-CLIP checkpoint (text encoder).
         init_gate: Initial fusion gate (image weight).
         baseline_mode: If True, gate = 1.0 (text ignored), DGTRS not loaded.
+        shared_encoder: If True, use single DINOv3 WEB for both branches.
         device: Torch device string.
     """
 
@@ -280,6 +287,7 @@ class AsymmetricEncoder(nn.Module):
         lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
         init_gate: float = 0.7,
         baseline_mode: bool = False,
+        shared_encoder: bool = True,
         mona_bottleneck: int = 64,
         lora_rank: int = 4,
         device: str = "cuda",
@@ -287,15 +295,25 @@ class AsymmetricEncoder(nn.Module):
         super().__init__()
         self.embed_dim = self.DINO_DIM
         self.baseline_mode = baseline_mode
+        self.shared_encoder = shared_encoder
         self.device = device
 
-        # Image encoders (frozen + MONA adapters).
+        # Image encoder(s) (frozen + MONA adapters).
+        if shared_encoder:
+            # Single DINOv3 WEB for both branches — saves ~4-5 GB VRAM.
+            self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path)
+            self._freeze(self.image_encoder)
+            inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck)
+            LOGGER.info("Shared encoder mode: single DINOv3 WEB for drone + satellite")
+        else:
+            # Separate encoders (asymmetric mode).
             self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path)
             self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path)
             self._freeze(self.drone_encoder)
             self._freeze(self.sat_encoder)
             inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck)
             inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck)
+            LOGGER.info("Asymmetric encoder mode: DINOv3 WEB (drone) + DINOv3 SAT (satellite)")
 
         # Text encoder — official DGTRS architecture (frozen + LoRA).
         if not baseline_mode:
@@ -324,10 +342,14 @@ class AsymmetricEncoder(nn.Module):
 
     def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
         """Encode drone images with MONA adapters. Returns [B, DINO_DIM]."""
+        if self.shared_encoder:
+            return self.image_encoder(images)
         return self.drone_encoder(images)
 
     def encode_satellite(self, images: torch.Tensor) -> torch.Tensor:
         """Encode satellite images with MONA adapters. Returns [B, DINO_DIM]."""
+        if self.shared_encoder:
+            return self.image_encoder(images)
         return self.sat_encoder(images)
 
     def encode_text_levels(
@@ -420,6 +442,7 @@ class AsymmetricEncoder(nn.Module):
         ckpt = {
             "model_state": self.state_dict(),
             "baseline_mode": self.baseline_mode,
+            "shared_encoder": self.shared_encoder,
             **extra,
         }
         tmp = path.with_suffix(path.suffix + ".tmp")
@@ -454,6 +477,7 @@ class AsymmetricEncoder(nn.Module):
             dino_sat_path=dino_sat_path,
             lrsclip_path=lrsclip_path,
             baseline_mode=ckpt.get("baseline_mode", False),
+            shared_encoder=ckpt.get("shared_encoder", True),
             device=device,
         )
         model.load_state_dict(ckpt["model_state"], strict=False)
@@ -464,11 +488,12 @@ class AsymmetricEncoder(nn.Module):
     def train(self, mode: bool = True) -> AsymmetricEncoder:
         """Override to keep frozen encoders in eval mode."""
         super().train(mode)
+        if self.shared_encoder:
+            self.image_encoder.eval()
+        else:
             self.drone_encoder.eval()
             self.sat_encoder.eval()
         if self.text_encoder is not None:
-            # Text encoder partially unfrozen — set to train mode
-            # but frozen layers won't update anyway.
             self.text_encoder.train(mode)
         return self
 

src/training/train_gtauav.py

@@ -73,6 +73,7 @@ class TrainConfigGTAUAV:
     lrsclip_path: str = _LRSCLIP
     init_gate: float = 0.7
     baseline_mode: bool = False
+    shared_encoder: bool = True  # single DINOv3 WEB for both branches (saves ~4-5 GB)
 
     # Training.
     resume_from: str | None = None  # path to checkpoint for resuming
@@ -341,13 +342,15 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         start_epoch = resume_ckpt.get("epoch", -1) + 1
     else:
         mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)"
-        LOGGER.info("Building model — %s", mode_str)
+        enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)"
+        LOGGER.info("Building model — %s, %s", mode_str, enc_str)
         model = AsymmetricEncoder(
             dino_web_path=cfg.dino_web_path,
             dino_sat_path=cfg.dino_sat_path,
             lrsclip_path=cfg.lrsclip_path,
             init_gate=cfg.init_gate,
             baseline_mode=cfg.baseline_mode,
+            shared_encoder=cfg.shared_encoder,
             device=cfg.device,
         ).to(cfg.device)