MONA only on last 12 blocks (skip early low-level blocks)

- Add last_n_blocks parameter to inject_mona_into_dinov3 (default=12)
- Blocks 0-11: pure frozen DINOv3 (low-level features, domain-agnostic)
- Blocks 12-23: MONA adapted (high-level semantic features)
- MONA params: 3.5M (was 6.85M, -49%)
- Total trainable: ~5.7M with text (was ~9.0M)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 21:51:52 +03:00
parent 4a0336e0ff
commit 916854f124

View File

@@ -149,24 +149,32 @@ def inject_mona_into_dinov3(
model: nn.Module,
bottleneck: int = 64,
dropout: float = 0.1,
last_n_blocks: int = 12,
) -> int:
"""Inject MONA adapters into a frozen DINOv3ViT model.
Adds two MonaAdapter modules per block (after attention, after MLP).
Only adapter parameters are trainable.
Only adapter parameters are trainable. Early blocks remain pure frozen
DINOv3 (low-level features don't need spatial adaptation).
Args:
model: DINOv3ViT model (already frozen).
bottleneck: MONA bottleneck dimension.
dropout: MONA dropout rate.
last_n_blocks: Number of last blocks to adapt (default 12 of 24).
Returns:
Number of trainable adapter parameters added.
"""
dim = model.embed_dim
n_adapters = 0
total_blocks = len(model.layer)
start_idx = total_blocks - last_n_blocks
for i, block in enumerate(model.layer):
if i < start_idx:
continue # Skip early blocks — frozen, no adaptation.
for block in model.layer:
# Create adapters.
block.mona_attn = MonaAdapter(dim, bottleneck, dropout)
block.mona_mlp = MonaAdapter(dim, bottleneck, dropout)
@@ -195,8 +203,8 @@ def inject_mona_into_dinov3(
n_params = sum(p.numel() for n, p in model.named_parameters() if "mona" in n)
LOGGER.info(
"🔧 MONA injected: %d adapters (%d per block), %s trainable params (bottleneck=%d)",
n_adapters, 2, f"{n_params:,}", bottleneck,
"🔧 MONA injected: %d adapters (blocks %d-%d of %d), %s trainable params (bottleneck=%d)",
n_adapters, start_idx, total_blocks - 1, total_blocks, f"{n_params:,}", bottleneck,
)
return n_params