Simplify model: shared DINOv3 WEB + MONA in last 12/24 blocks

Three related architecture changes, driven by a cost/simplicity trade-off:

1. **Shared encoder**: one DINOv3 LVD-1689M (WEB) processes both drone
   and satellite images. Previously asymmetric — separate WEB (drone) and
   SAT-493M (satellite) encoders. Saves ~303M frozen params and halves
   VRAM for the image tower. Expected to lose some satellite-domain
   inductive bias; MONA adapters pick up the slack.

2. **MONA in last 12/24 blocks**: adapters injected only in the top half
   of the ViT. The lowest 12 blocks keep their pretrained features
   untouched. Trainable MONA count drops from 14.0M (48 adapters × 2
   encoders) to 3.5M (24 adapters × 1 encoder).

3. **No DINO_SAT**: `nn_models/DINO_SAT` is no longer loaded by the
   default config. It stays on disk and the path param is kept for
   backward compat with asymmetric checkpoints.

Parameter counts (with text fusion + LoRA + gates):
  Before: 17.6M trainable / 733M total (2.35%)
  After:   7.06M trainable / 434M total (1.63%)

Also fixes a pre-existing resume bug: checkpoints now record
`shared_encoder`, `baseline_mode`, `mona_bottleneck`, `mona_last_n_blocks`
so `AsymmetricEncoder.load_checkpoint` can rebuild the right architecture.
Old checkpoints still load (missing keys fall back to asymmetric defaults
via `ckpt.get(..., <default>)`).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-24 16:26:17 +03:00
parent d98d853455
commit cb477f4b40
4 changed files with 32 additions and 21 deletions

View File

@@ -1,10 +1,13 @@
# Caption Quality Test for Cross-View Geo-Localization
## Архитектура системы (v3, 2026-04-21) — GTA-UAV эксперимент
## Архитектура системы (v3, 2026-04-24) — GTA-UAV эксперимент
```
Shared DINOv3 ViT-L/16 (LVD-1689M, frozen + MONA in last 12/24 blocks)
для обеих веток — drone и satellite кодируются одним encoder.
QUERY BRANCH (drone + L1/L2/L3 captions):
drone_img [B,3,256,256] --> DINOv3 ViT-L/16 LVD-1689M (frozen) --> d_img [B,1024]
drone_img [B,3,256,256] --> DINOv3 ViT-L/16 (shared) --> d_img [B,1024]
|
L1 --> DGTRS-CLIP (248 tok) --> z₁ [768] --\ |
L2 --> DGTRS-CLIP (248 tok) --> z₂ [768] ---+-- cat --> MLP(2304→1024→1024) --> d_txt [B,1024]
@@ -15,7 +18,7 @@ QUERY BRANCH (drone + L1/L2/L3 captions):
q̂ = q/‖q‖₂ --> query [B,1024]
GALLERY BRANCH (satellite + satellite captions):
sat_img [B,3,256,256] --> DINOv3 ViT-L/16 SAT-493M (frozen) --> s_img [B,1024]
sat_img [B,3,256,256] --> DINOv3 ViT-L/16 (shared, same weights) --> s_img [B,1024]
|
sat_L1 --> DGTRS-CLIP --> z₁ --\ |
sat_L2 --> DGTRS-CLIP --> z₂ ---+-- cat --> MLP (shared) --> s_txt [B,1024]
@@ -54,15 +57,16 @@ 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: 17.6M из 748M (2.35%)
- **MONA adapters** (2×DINOv3): 14.0M (2 per block × 24 × 2 encoders, bottleneck=64)
### Trainable parameters: 7.06M из 434M (1.63%)
- **MONA adapters** (shared DINOv3): 3.5M (2 per block × 12 last blocks, 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)
- DINOv3 x2 + DGTRS: frozen backbone weights
- DINOv3 (1 encoder) + DGTRS: frozen backbone weights
- **Без projection layers** — retrieval space = DINOv3 native 1024-dim
- **AMP:** frozen layers fp16, adapters + loss fp32
- **Примечание:** ранее была asymmetric setup (2×DINOv3 WEB+SAT, MONA во всех 24 блоках) с 17.6M trainable / 733M total. Упростили до shared + last-12 MONA.
### Optimizer & Scheduler
- **AdamW** с per-group LR: projections lr=1e-4, text encoder lr=1e-5
@@ -123,20 +127,14 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization.
## Backbones (v3)
### DINOv3 ViT-L/16 — Drone (web pretrained)
### DINOv3 ViT-L/16 — Shared (web pretrained)
- **Checkpoint:** `nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth`
- **Arch:** ViT-L/16, 24 layers, 16 heads, hidden=1024, MLP=4096, 303M params
- **Input:** 256x256, ImageNet normalization, patch=16 → 256 patches
- **Register tokens:** 4, RoPE theta=100.0
- **Status:** frozen
### DINOv3 ViT-L/16 — Satellite (sat pretrained)
- **Checkpoint:** `nn_models/DINO_SAT/model.safetensors`
- **HuggingFace:** `facebook/dinov3-vitl16-pretrain-sat493m`
- **Arch:** идентична DINO_WEB (ViT-L/16, hidden=1024, 303M params)
- **Input:** 256x256
- **Config:** `nn_models/DINO_SAT/config.json` — BROKEN (auth error), используем конфиг от DINO_WEB
- **Status:** frozen
- **MONA:** 24 адаптера в последних 12 блоках (blocks 12-23), bottleneck=64, 3.5M trainable
- **Status:** frozen кроме MONA
- **Примечание:** ранее asymmetric — использовался отдельно `nn_models/DINO_SAT/model.safetensors` (sat493m pretrain) для satellite ветки. Упростили до shared WEB-энкодера.
### DGTRS-CLIP ViT-L-14 (LRSCLIP) — Text encoder
- **Checkpoint:** `nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt`
@@ -306,10 +304,10 @@ python -m scripts.compare_runs \
### V3 (GTA-UAV, DINOv3 ViT-L/16, 256x256)
| Фаза | Оценка |
|------|--------|
| VRAM: 2x DINOv3-L + LRSCLIP + batch 64 | ~18-22 GB |
| VRAM: DINOv3-L (shared) + LRSCLIP + batch 64 | ~10-14 GB (было ~18-22 с 2× DINOv3) |
| GPU mem (smoke test, batch 4) | 3.1 GB |
| Batch size | 64 (default) |
| Total params | 733M (10.9M trainable, 1.49%) |
| Total params | 434M (7.06M trainable, 1.63%) — shared encoder + MONA в last 12/24 blocks |
### V2 (UAV-GeoLoc, GeoRSCLIP)
| Фаза | Время |