diff --git a/README.md b/README.md index 911d67f..5d4ab96 100644 --- a/README.md +++ b/README.md @@ -130,7 +130,22 @@ 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} ``` -MONA runs in **fp16** (AMP-native) with gradient checkpointing to save VRAM. +MONA runs in **bfloat16** with gradient checkpointing to save VRAM. +Applied only to the **last 12 blocks** (out of 24) — early blocks extract low-level +features (edges, textures) that are domain-agnostic and don't need spatial adaptation. + +**Why bfloat16, not fp16:** + +MONA's scaled LayerNorm uses `gamma` initialized at `1e-6` for near-identity output at start. +fp16 has min normal ~6.1e-5, so `1e-6` falls into the subnormal range where precision collapses, +causing NaN after a few blocks. bfloat16 has the same exponent range as fp32 (min ~1.2e-38), +so `1e-6` is safely representable. RTX 4090 supports bf16 natively with comparable throughput. + +| Precision | `1e-6` representable | MONA stable | VRAM (bs=48) | +|-----------|:---:|:---:|:---:| +| fp32 | yes | yes | 21.4 GB | +| **bfloat16** | **yes** | **yes** | **21.8 GB** | +| fp16 | subnormal (lossy) | **NaN** | — | **Why MONA over LoRA for DINOv3:** @@ -305,14 +320,14 @@ 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, fp16) | 6.85M | 6.85M | 2 per block × 24 blocks, bottleneck=64 | +| MONA adapters (shared, bf16) | 3.5M | 3.5M | 2 per block × last 12 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) | 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)** | **436M** | **9.0M (2.06%)** | retrieval dim = 512 | +| **Total (shared)** | **432M** | **5.6M (1.30%)** | retrieval dim = 512 | > **Asymmetric mode** (`--shared-encoder false`): uses separate DINOv3 WEB (drone) + DINOv3 SAT > (satellite) encoders with independent MONA adapters. Requires ~4-5 GB more VRAM. diff --git a/src/models/adapters.py b/src/models/adapters.py index 4727cf3..3d6d368 100644 --- a/src/models/adapters.py +++ b/src/models/adapters.py @@ -68,7 +68,11 @@ class MonaAdapter(nn.Module): self.gammax = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor, hw: tuple[int, int]) -> torch.Tensor: - """Apply MONA adapter in fp32 (gamma=1e-6 and small bottleneck underflow in fp16). + """Apply MONA adapter in bfloat16. + + bfloat16 has the same exponent range as fp32 (min ~1.2e-38), so gamma=1e-6 + is safe. fp16 would underflow (min normal ~6.1e-5). RTX 4090 supports bf16 + natively with comparable throughput to fp16. Args: x: Token features [B, N, D] where N includes CLS + register + patch tokens. @@ -78,10 +82,9 @@ class MonaAdapter(nn.Module): Adapted features [B, N, D] (same shape, residual connection). """ orig_dtype = x.dtype - with torch.amp.autocast("cuda", enabled=False): - x = x.float() + with torch.amp.autocast("cuda", dtype=torch.bfloat16): identity = x - # Scaled LayerNorm. + # Scaled LayerNorm (gamma=1e-6 safe in bf16). x = self.norm(x) * self.gamma + x * self.gammax x = self.down(x) # [B, N, bottleneck]