DINO-MX: Modular SSL Vision Framework

• DINO-MX is a modular self-supervised vision framework that integrates DINO, DINOv2, and DINOv3 into a single configuration-driven pipeline.
• It offers interchangeable components for backbone selection, augmentation policies, loss composition, and parameter-efficient adaptation like LoRA and layer freezing.
• The framework supports both natural and specialized image data with distributed training, interpretability tools, and label-guided augmentation for improved localization.

DINO-MX is a modular and extensible training framework for self-supervised vision learning that combines the core principles of DINO, DINOv2 and DINOv3 within a unified configuration-driven system. It is designed for transformer-based architectures, is fully compatible with the Hugging Face ecosystem, supports both natural and specialized data types including single- and multi-channel images, and integrates parameter-efficient adaptation, knowledge distillation, distributed training, interpretability tools, and a label-guided data augmentation method into a reproducible and scalable foundation for developing, adapting, and benchmarking self-supervised vision models across research and real-world applications (Gokmen et al., 3 Nov 2025).

1. Conceptual scope and design goals

DINO-MX is positioned as a response to the observation that existing self-supervised training pipelines for Vision Foundation Models are often inflexible, domain-specific, or computationally expensive. Its central design choice is to expose model construction, optimization, and distributed execution through configuration rather than through a tightly coupled implementation. Every training run is driven by two configuration files: one for accelerator and distributed settings, and one for SSL model and training settings (Gokmen et al., 3 Nov 2025).

The framework is explicitly modular. It is built around a “factory” of interchangeable components, so backbone choice, head selection, augmentation policy, parameter-efficient fine-tuning strategy, and distributed engine can be varied without rewriting the pipeline. A plausible implication is that DINO-MX is best understood less as a single algorithmic variant than as an experimental substrate for composing DINO-family methods under controlled and reproducible settings.

Its intended scope is broad. The framework is described as supporting both natural and specialized data types, including single- and multi-channel images. This matters in practice because the same infrastructure is used for RGB-oriented ViTs, single-channel medical imagery, and distillation from domain-specific teachers, rather than assuming one canonical image modality or one canonical pretraining corpus (Gokmen et al., 3 Nov 2025).

2. Component architecture and backbone integration

At the architectural level, DINO-MX exposes a set of core modules whose interfaces are standardized around transformer backbones and SSL heads.

Module	Function	Key details
Backbone Module	Loads ViTs	Hugging Face AutoConfig and AutoModel; standardized queries, keys, values, patch embeddings
Augmentation Module	Builds multi-view inputs	Global + local crops; RGB and single-channel medical presets; label-guided plugin
DINO / iBOT / Distillation Heads	Produces SSL targets	out_dim default 65 536; patch-masked auxiliary loss; frozen-teacher cross-entropy
PEFT Module	Reduces adaptation cost	LoRA in Q/V projection matrices; layer freezing
Loss Module	Composes objectives	Student–teacher DINO loss, iBOT loss, distillation loss, centering, sharpening
Distributed Engine	Scales execution	DDP or FSDP; 1–N GPUs; single/multi-node; fp16/bf16
Logging & Evaluation	Monitors and probes representations	Console loss/LR/GPU memory; linear-probe, k-NN, attention-map PCA, clustering

The backbone layer is explicitly tied to Hugging Face’s transformers library. Supported backbones include DINO-v1 ViTs such as facebook/dino-vitb8 and facebook/dino-vitb16, DINO-v2 ViTs such as facebook/dinov2-small, facebook/dinov2-base, and facebook/dinov2-large, and custom medical VFMs such as Prov-Gigapath. The wrapper around AutoConfig and AutoModel exposes a standardized interface for queries, keys, values, and patch-embedding extraction, which means that any future HF-compatible ViT can plug into DINO-MX by editing the YAML or JSON model_type: field rather than modifying code (Gokmen et al., 3 Nov 2025).

The augmentation module implements multi-crop training with domain-specific presets for RGB versus single-channel medical images. When pixel or box labels are available, a label-guided augmentation plugin forces additional crops around regions of interest. On the head side, the DINO head projects backbone outputs to a high-dimensional vector, the iBOT head adds a patch-masked reconstruction auxiliary loss, and the distillation head wraps a frozen teacher model that may be either a larger DINOv2 or a pathology-specific foundation model.

3. Objective functions and parameter-efficient adaptation

DINO-MX defines its training criterion as a weighted combination of DINO, iBOT, and distillation terms:

$$L_{\mathrm{total}} = L_{\mathrm{DINO}} + \alpha_{\mathrm{ibot}} L_{\mathrm{iBOT}} + \alpha_{\mathrm{distill}} L_{\mathrm{distill}}.$$

For the student–teacher DINO objective, student outputs $z_i^s$ and teacher outputs $z_j^t$ are transformed by sharpening and centering:

$$p_i^t = \operatorname{softmax}\bigl((z_i^t - c)/T_t\bigr), \qquad q_j^s = \operatorname{softmax}(z_j^s / T_s).$$

The cross-entropy term is

$$L_{\mathrm{DINO}} = - \sum_i \sum_j p_i^t \cdot \log q_j^s,$$

and the teacher center is updated as

$c \leftarrow m c + (1-m)\operatorname{mean}(z^t).$

The sharpening-and-centering operator is written componentwise as

$$\operatorname{softmax}_k(x) = \frac{\exp((x_k-c_k)/T)}{\sum_\ell \exp((x_\ell-c_\ell)/T)}.$$

The iBOT term is a masked-patch prediction loss with cross-entropy between masked patch classes, while the distillation term is

$$L_{\mathrm{distill}} = - \sum p^{\mathrm{teacher}} \cdot \log q^{\mathrm{student}}.$$

Three adaptation strategies are emphasized. Layer freezing freezes the first $N$ transformer blocks and leaves the remainder trainable. Low-Rank Adaptation injects trainable low-rank updates into the Q/V projection matrices according to

$$\Delta W = BA, \qquad B \in \mathbb{R}^{d\times r}, \quad A \in \mathbb{R}^{r\times k}, \quad r \ll \min(d,k).$$

Knowledge distillation adds student–teacher cross-entropy on the final DINO-head outputs. In the provided configuration examples, these strategies are exposed through settings such as freeze_backbone_layers, use_lora, lora_r, lora_alpha, lora_dropout, do_distillation, and a teacher specification such as distilled_model_type: 'facebook/dinov2-giant' (Gokmen et al., 3 Nov 2025).

These choices make the optimization layer compositional. A plausible implication is that DINO-MX treats PEFT, distillation, and SSL loss design as jointly tunable axes of experimentation rather than as mutually exclusive training recipes.

4. Configuration model and distributed execution

The framework abstracts distributed training through a single configuration switch between Distributed Data Parallel and Fully Sharded Data Parallel. It works on 1–N GPUs, in single-node or multi-node settings, and supports mixed precision in fp16 and bf16. The implementation is described as abstracting PyTorch Accelerate’s DDP versus FSDP through YAML configuration rather than through separate code paths (Gokmen et al., 3 Nov 2025).

The distinction between DDP and FSDP is treated operationally. DDP mirrors the full model on each GPU and performs gradient all-reduce. FSDP shards parameters, gradients, and optimizer states across GPUs on the fly. The corresponding trade-off is explicit: DDP is simpler but requires a full model replica per GPU, while FSDP reduces memory pressure at the cost of resharing overhead during forward and backward passes. Relevant control parameters include bucket_cap_mb, cpu_offload, gradient_accumulation_steps, and mixed-precision flags.

The configuration skeleton in the paper illustrates the broader training surface: out_dim: 65536, mask_ratio_min_max: [0.1, 0.5], global_crops_number: 2, global_crops_scale: [0.4, 1.0], local_crops_number: 8, local_crops_scale: [0.1, 0.4], global_batch_size: 64, max_iterations: 2000, lr: 1e-4, min_lr: 1e-5, and saveckp_freq: 250. These are presented as a training configuration skeleton rather than as universally optimal defaults. In encyclopedic terms, the notable point is that DINO-MX makes accelerator policy, crop policy, SSL head composition, and PEFT strategy first-class configuration objects rather than hidden implementation details (Gokmen et al., 3 Nov 2025).

5. Label-guided augmentation and interpretability

A distinctive extension in DINO-MX is the label-guided augmentation variant termed DINO-LG. Standard DINO samples random local crops, whereas DINO-LG also samples bounding-box or mask labels. For each annotated ROI, centers $z_i^s$0 are sampled within the ROI, after which local crops of size, for example, $z_i^s$1 are extracted. These guided crops receive the same medical-safe augmentations as standard local crops, specifically Gaussian blur, brightness, and noise addition (Gokmen et al., 3 Nov 2025).

The purpose of this mechanism is not to attach a separate detection or segmentation head, but to bias the crop distribution toward semantically relevant spatial regions during SSL training. The paper reports that this label-guided data augmentation improves attention-based localization without the need for extra detection or segmentation heads. This suggests that DINO-MX treats annotation not only as supervision for downstream evaluation but also as a way of shaping view generation.

Its interpretability stack is attention-centric. CLS-to-patch attention is extracted from each head as

$z_i^s$2

The CLS row $z_i^s$3 is flattened to length $z_i^s$4 and reshaped to $z_i^s$5; the 12 heads are stacked into a $z_i^s$6 tensor; PCA is then applied across the 12 channels to obtain the top-$z_i^s$7 components, often $z_i^s$8 or $z_i^s$9, for visualization. Thresholding or clustering on the PCA map is then used to localize ROIs. In this design, attention-map PCA and clustering serve as offline analysis tools that complement linear-probe and k-NN evaluation rather than replacing them (Gokmen et al., 3 Nov 2025).

6. Experimental regime, empirical results, and practical use

The reported experimental setup uses 2× NVIDIA A6000 GPUs with bf16 mixed precision, 2000 iterations, and global batch size 64, corresponding to 32 samples per GPU. The datasets listed for evaluation include MedMNIST v2 subsets—BloodMNIST, PathMNIST, DermaMNIST, and OrganAMNIST—together with a CT calcification dataset with pixel-level annotations for DINO-LG and a Prov-Gigapath pathology set for distillation. The two primary evaluation protocols are linear-probe, in which the backbone and DINO-MX head are frozen and a single linear layer is trained, and k-Nearest Neighbors with $z_j^t$0 on frozen embeddings (Gokmen et al., 3 Nov 2025).

Scenario	Setting	Reported result
MedMNIST evaluation	k-NN vs linear probe	k-NN beats linear probe by 10–25%
BloodMNIST example	Accuracy	0.97 vs 0.95
PEFT efficiency	LoRA + layer freezing	training time reduced by 30–40%; memory by 30–35%
CT localization	DINO-LG	AP = 0.796, recall = 0.707, localization score = 0.901
Pathology distillation	Prov-Gigapath → dinov2-small	+4% linear-probe accuracy on PathMNIST

These results characterize DINO-MX as a framework that trades on adaptation efficiency as well as representation quality. The gains from LoRA and layer freezing are reported with minimal accuracy loss, while the label-guided augmentation experiments indicate that localization-oriented behavior can be improved through crop policy rather than through a dedicated detector. The distillation result on PathMNIST further shows that the framework is designed to accommodate domain-specific teachers in addition to generic DINO-family backbones.

The reproducibility guidance is unusually explicit. Reported hyperparameter ranges include learning rate from $z_j^t$1 to $z_j^t$2 with linear warmup over the first 1000 iterations, weight decay from $z_j^t$3 to $z_j^t$4 depending on dataset size, teacher momentum from $z_j^t$5 to $z_j^t$6, LoRA rank $z_j^t$7 from $z_j^t$8 to $z_j^t$9 with $$p_i^t = \operatorname{softmax}\bigl((z_i^t - c)/T_t\bigr), \qquad q_j^s = \operatorname{softmax}(z_j^s / T_s).$$0 and dropout $$p_i^t = \operatorname{softmax}\bigl((z_i^t - c)/T_t\bigr), \qquad q_j^s = \operatorname{softmax}(z_j^s / T_s).$$1, and freeze depth from 0 to 12 layers. For resource-constrained settings, the documented recommendations are to use LoRA with FSDP on a single node, reduce the local crop count from 8 to 4 to save GPU memory, use mixed-precision bf16 on A100 or A6000 for 2× speed, and perform an early exit after 1000 iterations with partial evaluation if budget is tight. For new domains, the prescribed procedure is to change dataset_path, adjust crop scales to match image resolution, enable DINO-LG if labels are available, start with a pretrained facebook/dinov2-base, then enable PEFT with LoRA and freezing of the first 4 layers, monitor GPU memory, switch between DDP and FSDP via distribution.type, and evaluate embeddings with k-NN first because this is often faster than linear-probe training (Gokmen et al., 3 Nov 2025).