diff --git a/README.md b/README.md index 17087f5..911d67f 100644 --- a/README.md +++ b/README.md @@ -106,16 +106,17 @@ where `[z₁ ; z₂ ; z₃] ∈ ℝ^{B×2304}` is the concatenation of three 768 - `α_q, α_g` — separate learnable scalars in logit-space, init `σ(α) ≈ 0.7` - `σ` — sigmoid function -- `d_img, s_img ∈ ℝ^{B×1024}` — DINOv3+MONA image embeddings -- `d_txt, s_txt ∈ ℝ^{B×1024}` — fused text embeddings +- `d_img, s_img ∈ ℝ^{B×512}` — DINOv3+MONA → projection(1024→512) +- `d_txt, s_txt ∈ ℝ^{B×512}` — fused text embeddings (TextFusionMLP → 512) - For satellite images without captions: `s_txt = None → g = s_img` ### Adaptation methods | Method | Applied to | Where | Params | |--------|-----------|-------|--------| -| **MONA** (CVPR 2025) | DINOv3 ViT-L/16 (drone + sat) | After MSA and MLP in each of 24 blocks | 7.0M per encoder | +| **MONA** (CVPR 2025) | DINOv3 ViT-L/16 (shared) | After MSA and MLP in each of 24 blocks | 6.85M | | **LoRA** (rank=4) | DGTRS-CLIP text encoder | Q and V projections in all 12 blocks | 147K | +| **Projection** | After DINOv3 CLS output | Linear(1024→512) | 525K | **MONA adapter** (per block): @@ -129,13 +130,43 @@ where `x̂ = γ · LN(x) + γₓ · x` (scaled LayerNorm, `γ` init `10⁻⁶`, \text{MonaOp}(\mathbf{x}) = \frac{\text{DWConv}_{3 \times 3}(\mathbf{x}) + \text{DWConv}_{5 \times 5}(\mathbf{x}) + \text{DWConv}_{7 \times 7}(\mathbf{x})}{3} + \mathbf{x} ``` -**LoRA** (per attention layer): +MONA runs in **fp16** (AMP-native) with gradient checkpointing to save VRAM. + +**Why MONA over LoRA for DINOv3:** + +MONA uses multi-scale depthwise convolutions (3×3, 5×5, 7×7) that provide **spatial inductive bias** +critical for cross-view geo-localization. Drone images (oblique, 100-600m altitude) and satellite +images (nadir) exhibit a strong **geometric domain gap** — the same building looks spatially different +from each viewpoint. MONA's multi-scale spatial filters learn scale-invariant features to bridge +this gap, while LoRA (pure linear low-rank correction) would only handle style/distribution shifts. + +| | MONA | LoRA (on DINOv3) | +|---|---|---| +| Inductive bias | 2D spatial (knows about pixel neighbors) | None (linear correction) | +| Best for | Geometric domain gap (aerial↔satellite) | Style/distribution shift | +| Params | 6.85M (bottleneck=64) | ~0.3M (rank=4) | +| Compute | Heavier (192 conv ops per forward) | Light | +| CVGL fit | Strong (multi-scale spatial adaptation) | Weak (no spatial awareness) | + +**LoRA** (per DGTRS-CLIP attention layer): ```math \mathbf{Q}' = \mathbf{Q} + \frac{\alpha}{r} \cdot \mathbf{x} \mathbf{A}_Q^T \mathbf{B}_Q^T, \qquad \mathbf{V}' = \mathbf{V} + \frac{\alpha}{r} \cdot \mathbf{x} \mathbf{A}_V^T \mathbf{B}_V^T ``` -where `A ∈ ℝ^{r×d}`, `B ∈ ℝ^{d×r}`, `r = 4` +where `A ∈ ℝ^{r×d}`, `B ∈ ℝ^{d×r}`, `r = 4`. LoRA is appropriate for the text encoder since +text has no spatial structure — the adaptation needed is purely semantic/distributional. + +### Projection head (DINOv3 1024 → 512) + +```math +\mathbf{h} = \text{Linear}_{1024 \to 512}\!\bigl(\text{CLS}_{\text{DINOv3}}\bigr) +``` + +Reduces the retrieval space from DINOv3 native 1024-dim to 512-dim. Benefits: +- Smaller similarity matrix in InfoNCE (`B×B` at 512 vs 1024) +- TextFusionMLP outputs 512 instead of 1024 (fewer params: 1.5M vs 3.4M) +- Shared projection for both drone and satellite branches ### Pair formation and negative sampling @@ -274,17 +305,17 @@ Mixed precision: AMP fp16 for model forward, fp32 for loss | Component | Params | Trainable | Notes | |-----------|--------|-----------|-------| | DINOv3 ViT-L/16 WEB (shared) | 303M | frozen | single encoder for drone + satellite | -| MONA adapters (shared) | 7.0M | 7.0M | 2 per block × 24 blocks, bottleneck=64 | +| MONA adapters (shared, fp16) | 6.85M | 6.85M | 2 per block × 24 blocks, bottleneck=64 | +| Image projection | 525K | 525K | Linear(1024→512) after DINOv3 CLS | | 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) | +| TextFusionMLP (shared) | 1.5M | 1.5M | Linear(2304,512) + GELU + Linear(512,512) | | GatedFusion α_q + α_g | 2 | 2 | separate gate scalars | | logit_scale | 1 | 1 | learnable temperature | -| **Total (shared)** | **438M** | **10.6M (2.42%)** | retrieval dim = 1024 | +| **Total (shared)** | **436M** | **9.0M (2.06%)** | retrieval dim = 512 | > **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. +> (satellite) encoders with independent MONA adapters. Requires ~4-5 GB more VRAM. ## Experiments