lmd-vitl-p8-simmim-scratch

A 3D ViT-L encoder for volume microscopy (electron and expansion microscopy) at 8 nm. It has the DINOv3 ViT-L architecture with 8 × 8 × 8-voxel patches and was pretrained from random initialization for 100k steps with SimMIM on LMD, a corpus of 243 microscopy volumes.

Use

import json
import sys
from pathlib import Path

import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file

sys.path.insert(0, "mia-train/src")  # git clone https://github.com/AI-HHMI/mia-train && git -C mia-train checkout 0fe3872
from models.dinov3_vit3d import DinoVisionTransformer3D

root = Path(snapshot_download("ai-hhmi/lmd-vitl-p8-simmim-scratch"))
config = json.loads((root / "config.json").read_text())
encoder = DinoVisionTransformer3D(**config["model"]["kwargs"])
encoder.load_state_dict(load_file(root / "model.safetensors"))

device = "cuda" if torch.cuda.is_available() else "cpu"
encoder = encoder.to(device).eval()

# (batch, channel, x, y, z): 8 nm voxels, intensities in [0, 1], every side a multiple of 8
x = torch.rand(1, 1, 128, 128, 128, device=device)
with torch.no_grad(), torch.autocast(device, dtype=torch.bfloat16, enabled=device == "cuda"):
    tokens, grid = encoder.patch_features(x)  # (1, 4096, 1024) final-layer patch tokens on a (16, 16, 16) grid
    maps = encoder.get_intermediate_layers(x, n=[7, 15, 23], reshape=True)  # 3 x (1, 1024, 16, 16, 16)
print(tuple(tokens.shape), grid, [tuple(m.shape) for m in maps])

Requirements: PyTorch (we used 2.13), NumPy, safetensors, huggingface_hub, and the model code from mia-train at commit 0fe3872. Put mia-train/src on the path as above rather than installing mia-train. Its top-level packages have generic names (models, layers, utils, distributed, data), and they shadow packages of the same name, such as Dask's distributed.

• encoder.patch_features(x) returns the final layer's normalised patch tokens and their grid. encoder.forward_features(x) also returns the CLS token and the 4 storage tokens. encoder.get_intermediate_layers(x, n=[...], reshape=True) returns feature maps from any blocks.
• The encoder is a plain torch.nn.Module, so you can fine-tune it in your own training loop. Our affinity models (table above) read blocks [7, 15, 23] through a UNETR decoder.

Input

• One channel, with axes ordered (x, y, z). Arrays stored (z, y, x), as OME-Zarr and N5 volumes usually are, need array.transpose(2, 1, 0).
• 8 nm isotropic voxels. Resample other data first.
• Intensities in [0, 1], scaled per volume as in training: 8-bit data divided by 255; 16-bit and float data clamped to the volume's 0.1st and 99.9th percentiles, then rescaled to [0, 1].
• Every side a multiple of 8 voxels. Training used 256³ windows. The model runs at other sizes, but we did not test them.

Model details

• DINOv3 ViT-L: 24 blocks, width 1024, 16 heads, MLP ratio 4, LayerScale, 4 storage tokens and a masked key bias. A 3D rotary position embedding gives each axis a third of the rotary channels; there is no position-embedding table. 303 M parameters.
• model.safetensors: the encoder's weights in fp32 (1.21 GB). It leaves out SimMIM's pixel-prediction head and the optimizer state.
• config.json: the constructor arguments (model.kwargs), the input conventions and the source training run.

Training

• Data. LMD, 243 volumes of 3D microscopy from 12 species: 177 electron microscopy volumes (fly, mouse, zebrafish, human, C. elegans, zebra finch and others) and 66 expansion microscopy volumes (fly, mouse, Danionella). We resampled all of them to 8 nm isotropic voxels, sampled volumes in proportion to size^0.4, and drew random 256³ crops.
• Objective. SimMIM: 90% of the input is hidden in 32³-voxel blocks, a linear head predicts the hidden voxels, and the loss is the L1 error on them.
• Initialization. Random.
• Optimization. 100k steps at a global batch of 32 on 8 GPUs. AdamW (β = 0.9, 0.95; weight decay 0.05, none on norms and biases). Learning rate 3e-4 after 5k warm-up steps, decaying linearly to 3e-7. Gradient clipping at 1.0, bf16 autocast. Augmentation: in-plane rotations only.

Evaluation

We did not score the encoder on its own. Fine-tuned for affinities, it gives ai-hhmi/lmd-vitl-p8-affinity-scratch, whose neuron segmentation scores are on that page and on our leaderboard.

License

DINOv3 License (LICENSE.md). This model holds no Meta weights, but its architecture code is a port of Meta's DINOv3 reference implementation, which that agreement covers, so we release it under the same agreement. Publications that use the model must acknowledge DINOv3.

References

• O. Siméoni et al. DINOv3. arXiv:2508.10104, 2025.
• Z. Xie et al. SimMIM: A Simple Framework for Masked Image Modeling. CVPR 2022.