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, model: nn.Module,
bottleneck: int = 64, bottleneck: int = 64,
dropout: float = 0.1, dropout: float = 0.1,
last_n_blocks: int = 12,
) -> int: ) -> int:
"""Inject MONA adapters into a frozen DINOv3ViT model. """Inject MONA adapters into a frozen DINOv3ViT model.
Adds two MonaAdapter modules per block (after attention, after MLP). 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: Args:
model: DINOv3ViT model (already frozen). model: DINOv3ViT model (already frozen).
bottleneck: MONA bottleneck dimension. bottleneck: MONA bottleneck dimension.
dropout: MONA dropout rate. dropout: MONA dropout rate.
last_n_blocks: Number of last blocks to adapt (default 12 of 24).
Returns: Returns:
Number of trainable adapter parameters added. Number of trainable adapter parameters added.
""" """
dim = model.embed_dim dim = model.embed_dim
n_adapters = 0 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. # Create adapters.
block.mona_attn = MonaAdapter(dim, bottleneck, dropout) block.mona_attn = MonaAdapter(dim, bottleneck, dropout)
block.mona_mlp = 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) n_params = sum(p.numel() for n, p in model.named_parameters() if "mona" in n)
LOGGER.info( LOGGER.info(
"🔧 MONA injected: %d adapters (%d per block), %s trainable params (bottleneck=%d)", "🔧 MONA injected: %d adapters (blocks %d-%d of %d), %s trainable params (bottleneck=%d)",
n_adapters, 2, f"{n_params:,}", bottleneck, n_adapters, start_idx, total_blocks - 1, total_blocks, f"{n_params:,}", bottleneck,
) )
return n_params return n_params