diff --git a/CLAUDE.md b/CLAUDE.md index c4400e5..47d29b4 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -50,14 +50,15 @@ BASELINE: σ(α_q)=σ(α_g)=1.0, text disabled, DGTRS not loaded - Transformer: sequence-first (LND), nn.MultiheadAttention, 12 layers - Tokenizer: BPE SimpleTokenizer (248 tokens, vocab 49408) -### Trainable parameters: 11.1M из 734M (1.51%) -- TextFusionMLP (shared): Linear(2304,1024)+GELU+Linear(1024,1024) = ~3.5M -- gate α_q: 1 scalar (query branch) -- gate α_g: 1 scalar (gallery branch) +### Trainable parameters: 17.6M из 748M (2.35%) +- **MONA adapters** (2×DINOv3): 14.0M (2 per block × 24 × 2 encoders, bottleneck=64) +- **LoRA** (DGTRS-CLIP): 147K (Q+V, rank=4, 12 blocks) +- TextFusionMLP (shared): Linear(2304,1024)+GELU+Linear(1024,1024) = ~3.4M +- gate α_q + α_g: 2 scalars - logit_scale: 1 scalar (learnable temperature) -- DGTRS partial unfreeze (last resblock + ln_final + text_projection): ~7.6M -- DINOv3 x2 (303M each): frozen +- DINOv3 x2 + DGTRS: frozen backbone weights - **Без projection layers** — retrieval space = DINOv3 native 1024-dim +- **AMP:** frozen layers fp16, adapters + loss fp32 ### Optimizer & Scheduler - **AdamW** с per-group LR: projections lr=1e-4, text encoder lr=1e-5 diff --git a/README.md b/README.md index 054a7b8..c2e5a9b 100644 --- a/README.md +++ b/README.md @@ -60,52 +60,64 @@ Each drone image has a VLM-generated caption (Qwen3-VL) split into 3 levels: | **L3** | Fingerprint | P3: unique landmarks and spatial signature for matching | 20–50 tokens | 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). +(248-token context via KPS positional embedding, 768-dim output) +adapted with **LoRA** (rank=4 on Q/V in all 12 blocks). -**Text fusion (shared MLP for both branches):** +### Text fusion (shared MLP for both branches) -``` -z_text = MLP( [z₁ ; z₂ ; z₃] ) +$$\mathbf{z}_{\text{text}} = \text{MLP}\bigl([\mathbf{z}_1 \;;\; \mathbf{z}_2 \;;\; \mathbf{z}_3]\bigr)$$ -where [z₁ ; z₂ ; z₃] ∈ ℝ^(B×2304) — concatenation of three 768-dim embeddings - MLP: Linear(2304, 1024) → GELU → Linear(1024, 1024) - z_text ∈ ℝ^(B×1024) -``` +where $[\mathbf{z}_1 ; \mathbf{z}_2 ; \mathbf{z}_3] \in \mathbb{R}^{B \times 2304}$ is the concatenation of three 768-dim DGTRS-CLIP embeddings, and -**Gated fusion (separate gates for query and gallery):** +$$\text{MLP}: \text{Linear}(2304, 1024) \to \text{GELU} \to \text{Linear}(1024, 1024), \quad \mathbf{z}_{\text{text}} \in \mathbb{R}^{B \times 1024}$$ -``` -q = σ(α_q) · d_img + (1 − σ(α_q)) · d_txt (query branch) -g = σ(α_g) · s_img + (1 − σ(α_g)) · s_txt (gallery branch) +### Gated fusion (separate gates for query and gallery) -where α_q, α_g — separate learnable scalars in logit-space (init: σ(α) ≈ 0.7) - σ — sigmoid function - d_img, s_img — DINOv3 image embeddings [B, 1024] - d_txt, s_txt — fused text embeddings [B, 1024] +$$\mathbf{q} = \sigma(\alpha_q) \cdot \mathbf{d}_{\text{img}} + \bigl(1 - \sigma(\alpha_q)\bigr) \cdot \mathbf{d}_{\text{txt}} \qquad \text{(query branch)}$$ -For satellite images without captions: s_txt = None → g = s_img -``` +$$\mathbf{g} = \sigma(\alpha_g) \cdot \mathbf{s}_{\text{img}} + \bigl(1 - \sigma(\alpha_g)\bigr) \cdot \mathbf{s}_{\text{txt}} \qquad \text{(gallery branch)}$$ + +- $\alpha_q, \alpha_g$ — separate learnable scalars in logit-space, init $\sigma(\alpha) \approx 0.7$ +- $\sigma$ — sigmoid function +- $\mathbf{d}_{\text{img}}, \mathbf{s}_{\text{img}} \in \mathbb{R}^{B \times 1024}$ — DINOv3+MONA image embeddings +- $\mathbf{d}_{\text{txt}}, \mathbf{s}_{\text{txt}} \in \mathbb{R}^{B \times 1024}$ — fused text embeddings +- For satellite images without captions: $\mathbf{s}_{\text{txt}} = \text{None} \Rightarrow \mathbf{g} = \mathbf{s}_{\text{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 | +| **LoRA** (rank=4) | DGTRS-CLIP text encoder | Q and V projections in all 12 blocks | 147K | + +**MONA adapter** (per block): + +$$\mathbf{x} \leftarrow \mathbf{x} + \text{Up}_{64 \to 1024}\!\Bigl(\text{GELU}\bigl(\text{MonaOp}\bigl(\text{Down}_{1024 \to 64}(\hat{\mathbf{x}})\bigr)\bigr)\Bigr)$$ + +where $\hat{\mathbf{x}} = \gamma \cdot \text{LN}(\mathbf{x}) + \gamma_x \cdot \mathbf{x}$ (scaled LayerNorm, $\gamma$ init $10^{-6}$, $\gamma_x$ init $1$) + +$$\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): + +$$\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 $\mathbf{A} \in \mathbb{R}^{r \times d}$, $\mathbf{B} \in \mathbb{R}^{d \times r}$, $r = 4$ ### Loss function Symmetric InfoNCE with learnable temperature (CLIP-style `logit_scale`): -``` -L = w_q2g · L_q→g + w_g2q · L_g→q +$$\mathcal{L} = w_{q \to g} \cdot \mathcal{L}_{q \to g} + w_{g \to q} \cdot \mathcal{L}_{g \to q}$$ -L_q→g = CrossEntropy( q̂ · ĝᵀ / τ, targets ) -L_g→q = CrossEntropy( ĝ · q̂ᵀ / τ, targets ) +$$\mathcal{L}_{q \to g} = \text{CrossEntropy}\!\left(\frac{\hat{\mathbf{q}} \cdot \hat{\mathbf{g}}^T}{\tau},\; \text{targets}\right), \qquad \mathcal{L}_{g \to q} = \text{CrossEntropy}\!\left(\frac{\hat{\mathbf{g}} \cdot \hat{\mathbf{q}}^T}{\tau},\; \text{targets}\right)$$ -where τ = 1 / exp(logit_scale), logit_scale — learnable scalar - τ ∈ [0.01, 0.5] (clamped) - τ_init = 0.07 - w_q2g = 0.6, w_g2q = 0.4 - targets = [0, 1, 2, ..., B−1] (positives on diagonal) - label_smoothing = 0.1 -``` +- $\tau = 1 / \exp(\text{logit\_scale})$ — learnable scalar, clamped $\tau \in [0.01, 0.5]$, init $\tau_0 = 0.07$ +- $w_{q \to g} = 0.6$, $w_{g \to q} = 0.4$ +- $\text{targets} = [0, 1, 2, \ldots, B-1]$ — positives on diagonal +- label smoothing $= 0.1$ -Loss is computed in **fp32** (outside AMP autocast) to prevent gradient overflow -in the learnable temperature. +Loss and adapters run in **fp32** (AMP autocast disabled) to prevent gradient overflow. ### Metrics @@ -120,9 +132,9 @@ Reported: R@1, R@5, R@10 for both q→g and g→q directions. ``` Optimizer: AdamW - - TextFusionMLP, gate α_q, gate α_g, logit_scale: + - MONA adapters, TextFusionMLP, gate α_q, gate α_g, logit_scale: lr = 1e-4, weight_decay = 1e-4 - - DGTRS text encoder (last resblock + ln_final + text_projection): + - LoRA adapters (DGTRS-CLIP text encoder): lr = 1e-5 (10× lower, --text-lr-factor 0.1) Scheduler: Linear warmup (2 epochs) + cosine annealing @@ -151,14 +163,16 @@ Mixed precision: AMP fp16 for model forward, fp32 for loss | Component | Params | Trainable | Notes | |-----------|--------|-----------|-------| -| 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 | -| 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 | +| DINOv3 ViT-L/16 LVD (drone) | 303M | frozen | backbone weights frozen | +| MONA adapters (drone) | 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** | **734M** | **11.1M (1.51%)** | retrieval dim = 1024 | +| **Total** | **748M** | **17.6M (2.35%)** | retrieval dim = 1024 | ## Experiments