Remove projections (1024 native), add satellite text, dual GatedFusion
Architecture changes: - Removed proj_drone/proj_sat (1024→512): retrieval space is now DINOv3 native 1024-dim, no information loss from projection - TextFusionMLP: 2304→1024→1024 (was 2304→768→512), shared between branches - Gallery branch now uses satellite captions (L1/L2/L3) via shared TextFusionMLP - Two separate GatedFusion gates: α_q (query) and α_g (gallery) - For sat images without captions (~57%): gate passes image features through Dataset changes: - GTAUAVDataset now loads satellite captions from caption index - collate_gtauav_batch includes sat_caption_l1/l2/l3 Training loop: - Passes satellite captions to model forward - Logs both gate_q and gate_g values 11.1M trainable / 734M total (1.51%) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
64
README.md
64
README.md
@@ -14,35 +14,39 @@ text fusion for drone-to-satellite image retrieval.
|
||||
┌──────────────────────────── QUERY BRANCH ────────────────────────────┐
|
||||
│ │
|
||||
│ drone_img ──► DINOv3 ViT-L/16 LVD ──► CLS token │
|
||||
│ [B,3,256,256] (frozen, 303M) [B,1024] │
|
||||
│ │ │
|
||||
│ proj_drone: Linear(1024,512) │
|
||||
│ │ │
|
||||
│ d_img [B,512] │
|
||||
│ [B,3,256,256] (frozen, 303M) d_img [B,1024] │
|
||||
│ │ │
|
||||
│ L1 (overview) ──► DGTRS-CLIP ──► z₁ [B,768] ─┐ │
|
||||
│ L2 (full desc) ──► DGTRS-CLIP ──► z₂ [B,768] ─┼─ cat ──► [B,2304]│
|
||||
│ L3 (fingerprint) ──► DGTRS-CLIP ──► z₃ [B,768] ─┘ │ │
|
||||
│ (248 tokens, KPS pos. emb.) MLP(2304→768→512) │
|
||||
│ │ │
|
||||
│ d_txt [B,512] │
|
||||
│ │ │
|
||||
│ q = σ(α)·d_img + (1−σ(α))·d_txt GatedFusion │
|
||||
│ (248 tokens, KPS pos. emb.) MLP(2304→1024→1024) │
|
||||
│ │ │
|
||||
│ d_txt [B,1024] │
|
||||
│ │ │
|
||||
│ q = σ(α_q)·d_img + (1−σ(α_q))·d_txt GatedFusion_q │
|
||||
│ │ │
|
||||
│ q̂ = q / ‖q‖₂ ──► query [B,512] │
|
||||
│ q̂ = q / ‖q‖₂ ──► query [B,1024] │
|
||||
└───────────────────────────────────────────────────────────────────────┘
|
||||
|
||||
┌──────────────────────────── GALLERY BRANCH ──────────────────────────┐
|
||||
│ │
|
||||
│ sat_img ──► DINOv3 ViT-L/16 SAT ──► CLS token │
|
||||
│ [B,3,256,256] (frozen, 303M) [B,1024] │
|
||||
│ [B,3,256,256] (frozen, 303M) s_img [B,1024] │
|
||||
│ │ │
|
||||
│ proj_sat: Linear(1024,512) │
|
||||
│ sat_L1 ──► DGTRS-CLIP ──► z₁ [768] ─┐ │
|
||||
│ sat_L2 ──► DGTRS-CLIP ──► z₂ [768] ─┼─ cat ──► MLP ──► s_txt [1024]│
|
||||
│ sat_L3 ──► DGTRS-CLIP ──► z₃ [768] ─┘ (shared MLP) │
|
||||
│ │ │
|
||||
│ ĝ = g / ‖g‖₂ ──► gallery [B,512] │
|
||||
│ g = σ(α_g)·s_img + (1−σ(α_g))·s_txt GatedFusion_g │
|
||||
│ │ │
|
||||
│ ĝ = g / ‖g‖₂ ──► gallery [B,1024]│
|
||||
└───────────────────────────────────────────────────────────────────────┘
|
||||
|
||||
BASELINE: σ(α) = 1.0 → q = d_img (text branch disabled, DGTRS not loaded)
|
||||
Retrieval space: 1024-dim (DINOv3 native, no projection layers)
|
||||
TextFusionMLP shared between query and gallery branches
|
||||
For sat images without captions: s_txt=None → g = s_img (gate passthrough)
|
||||
|
||||
BASELINE: σ(α) = 1.0 for both branches (text disabled, DGTRS not loaded)
|
||||
```
|
||||
|
||||
### Text hierarchy (L1 / L2 / L3)
|
||||
@@ -58,25 +62,28 @@ Each drone image has a VLM-generated caption (Qwen3-VL) split into 3 levels:
|
||||
All three levels are encoded by a **single DGTRS-CLIP ViT-L-14** text encoder
|
||||
(248-token context via KPS positional embedding, 768-dim output).
|
||||
|
||||
**Text fusion:**
|
||||
**Text fusion (shared MLP for both branches):**
|
||||
|
||||
```
|
||||
z_text = MLP( [z₁ ; z₂ ; z₃] )
|
||||
|
||||
where [z₁ ; z₂ ; z₃] ∈ ℝ^(B×2304) — concatenation of three 768-dim embeddings
|
||||
MLP: Linear(2304, 768) → GELU → Linear(768, 512)
|
||||
z_text ∈ ℝ^(B×512)
|
||||
MLP: Linear(2304, 1024) → GELU → Linear(1024, 1024)
|
||||
z_text ∈ ℝ^(B×1024)
|
||||
```
|
||||
|
||||
**Gated fusion:**
|
||||
**Gated fusion (separate gates for query and gallery):**
|
||||
|
||||
```
|
||||
q = σ(α) · d_img + (1 − σ(α)) · d_txt
|
||||
q = σ(α_q) · d_img + (1 − σ(α_q)) · d_txt (query branch)
|
||||
g = σ(α_g) · s_img + (1 − σ(α_g)) · s_txt (gallery branch)
|
||||
|
||||
where α — learnable scalar in logit-space (init: σ(α) ≈ 0.7)
|
||||
where α_q, α_g — separate learnable scalars in logit-space (init: σ(α) ≈ 0.7)
|
||||
σ — sigmoid function
|
||||
d_img — projected drone image embedding [B, 512]
|
||||
d_txt — fused text embedding [B, 512]
|
||||
d_img, s_img — DINOv3 image embeddings [B, 1024]
|
||||
d_txt, s_txt — fused text embeddings [B, 1024]
|
||||
|
||||
For satellite images without captions: s_txt = None → g = s_img
|
||||
```
|
||||
|
||||
### Loss function
|
||||
@@ -113,7 +120,7 @@ Reported: R@1, R@5, R@10 for both q→g and g→q directions.
|
||||
|
||||
```
|
||||
Optimizer: AdamW
|
||||
- Projection heads (proj_drone, proj_sat, TextFusionMLP, gate α, logit_scale):
|
||||
- TextFusionMLP, gate α_q, gate α_g, logit_scale:
|
||||
lr = 1e-4, weight_decay = 1e-4
|
||||
- DGTRS text encoder (last resblock + ln_final + text_projection):
|
||||
lr = 1e-5 (10× lower, --text-lr-factor 0.1)
|
||||
@@ -147,12 +154,11 @@ Mixed precision: AMP fp16 for model forward, fp32 for loss
|
||||
| DINOv3 ViT-L/16 LVD (drone) | 303M | 0 | frozen |
|
||||
| DINOv3 ViT-L/16 SAT (satellite) | 303M | 0 | frozen |
|
||||
| DGTRS-CLIP ViT-L-14 (text) | 124M | ~7.6M | last block + ln_final + text_projection |
|
||||
| proj_drone | 524K | 524K | Linear(1024, 512) |
|
||||
| proj_sat | 524K | 524K | Linear(1024, 512) |
|
||||
| TextFusionMLP | 2.2M | 2.2M | Linear(2304,768) + GELU + Linear(768,512) |
|
||||
| GatedFusion α | 1 | 1 | scalar |
|
||||
| TextFusionMLP (shared) | 3.5M | 3.5M | Linear(2304,1024) + GELU + Linear(1024,1024) |
|
||||
| GatedFusion α_q | 1 | 1 | query gate scalar |
|
||||
| GatedFusion α_g | 1 | 1 | gallery gate scalar |
|
||||
| logit_scale | 1 | 1 | learnable temperature |
|
||||
| **Total** | **733M** | **10.9M (1.49%)** | |
|
||||
| **Total** | **734M** | **11.1M (1.51%)** | retrieval dim = 1024 |
|
||||
|
||||
## Experiments
|
||||
|
||||
|
||||
Reference in New Issue
Block a user