MONA bfloat16: safe low-precision (gamma=1e-6 needs bf16 exponent range)

- Switch MONA from fp32 to bfloat16 (same exponent range as fp32, no underflow)
- fp16 causes NaN: gamma=1e-6 falls into subnormal range (min normal ~6.1e-5)
- bf16 min normal ~1.2e-38, so 1e-6 is safe
- RTX 4090 supports bf16 natively
- Document bf16 vs fp16 vs fp32 comparison in README
- Update model summary: 3.5M MONA (last 12 blocks), 5.6M total trainable

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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.

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]