lmd-vitl-p8-affinity-lvd

Predicts voxel affinities for neuron segmentation in 3D electron and expansion microscopy at 8 nm. The encoder is ai-hhmi/lmd-vitl-p8-simmim-lvd (DINOv3 LVD-1689M, then SimMIM on LMD). We fine-tuned it end to end with a UNETR decoder for 100k steps on 41 labelled volumes. This file holds the step-100,000 checkpoint.

Model	Starts from	Training
ai-hhmi/lmd-vitl-p8-simmim-lvd	DINOv3 ViT-L/16, LVD-1689M	SimMIM on LMD, 100k steps
ai-hhmi/lmd-vitl-p8-simmim-scratch	random	SimMIM on LMD, 100k steps
ai-hhmi/lmd-vitl-p8-affinity-lvd (this one)	lmd-vitl-p8-simmim-lvd	affinities, 100k steps
ai-hhmi/lmd-vitl-p8-affinity-scratch	lmd-vitl-p8-simmim-scratch	affinities, 75k steps
ai-hhmi/lmd-vitl-p8-cellmap-lvd	lmd-vitl-p8-simmim-lvd	CellMap organelle classes, 20k steps

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 layers.common.dense_heads import UNETRHead
from models.dinov3_vit3d import DinoVisionTransformer3D


class AffinityNet(torch.nn.Module):
    """(B, 1, X, Y, Z) image in [0, 1] -> (B, 6, X, Y, Z) affinity logits."""

    def __init__(self, config):
        super().__init__()
        self.encoder = DinoVisionTransformer3D(**config["model"]["kwargs"])
        self.decoder_out = UNETRHead(**config["head"]["kwargs"])
        self.skip_layers = config["skip_layers"]

    def forward(self, x):
        maps = self.encoder.get_intermediate_layers(x, n=self.skip_layers, reshape=True)
        return self.decoder_out(list(maps), x)


root = Path(snapshot_download("ai-hhmi/lmd-vitl-p8-affinity-lvd"))
config = json.loads((root / "config.json").read_text())
net = AffinityNet(config)
net.load_state_dict(load_file(root / "model.safetensors"))

device = "cuda" if torch.cuda.is_available() else "cpu"
net = net.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"):
    logits = net(x)  # (1, 6, 128, 128, 128), one channel per offset in config["offsets"]
affinities = torch.sigmoid(logits.float())  # P(voxel and voxel + offset belong to the same object)
print(tuple(affinities.shape))

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.

• Outputs. config["offsets"] gives each channel's offset in voxels along (x, y, z): (1, 0, 0), (0, 1, 0), (0, 0, 1), (10, 0, 0), (0, 10, 0), (0, 0, 10). Channel c at voxel v predicts whether v and v + offset_c belong to the same object. Background is affine to nothing.
• Weights. The file holds the encoder under encoder. and the decoder under decoder_out., so net.encoder can also start a different head.
• Inside mia-train. To use the model with mia-train's own trainer or predict.py, build AffinitySegmentation from config["algorithm"] and load the file with safetensors.torch.load_model. That function copes with the encoder's two names in that class. This route also needs miao-io and connected-components-3d.

Whole volumes

To score our benchmark, we predicted 256³ windows at a stride of 128 voxels under bf16 autocast. We mapped each window's logits to sigmoid(0.2 * logit) and averaged overlapping windows, with weights that rise linearly from a window's faces to its centre. We then ran a mutex watershed on those affinities and removed segments under 20,000 voxels, using mia-train's predict.py and mia-evals.

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 and our evaluation used 256³ windows. The model runs at other sizes, but we did not test them.

Model details

• Encoder: as in lmd-vitl-p8-simmim-lvd (DINOv3 ViT-L, 8³-voxel patches, 303 M parameters).
• Decoder: UNETR, 1.85 M parameters. It reads blocks 8, 16 and 24 (skip_layers = [7, 15, 23], 0-based) and the raw image. It upsamples three times by 2, with 64, 32 and 16 channels at strides 4, 2 and 1.
• model.safetensors: fp32, 1.22 GB.
• config.json: the constructor arguments of the encoder (model) and the decoder (head), skip_layers, offsets, mia-train's algorithm settings, the input conventions and the source training run.

Training

• Initialization. Encoder from lmd-vitl-p8-simmim-lvd (step 100k); decoder random.
• Data. 41 labelled volumes at 8 nm. 8 come from LMD: 6 electron microscopy volumes (mouse cortex, zebrafish, and fly, including CREMI A and B) and 2 expansion microscopy volumes (mouse dentate gyrus and fly). The other 33 are the zebra finch training cubes of the LSD release. Training ignored unannotated voxels. It read labels 10 voxels beyond each crop, so affinities that cross a crop's faces still had targets.
• Loss. Binary cross-entropy on the 6 affinity logits.
• Optimization. 100k steps at a global batch of 8 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, intensity scaling and shifts, noise, and dropped or shifted sections.

Evaluation

Mean panoptic quality (PQ) on the test split of our neuron segmentation benchmark: 7 held-out electron and expansion microscopy volumes at 8 nm. Segmentation: mutex watershed, with the minimum segment size (20,000 voxels) fitted on a separate fit split.

Benchmark	Test PQ
lmd_ssl_v2_neuron_instance_crop512 (centred crops of at most 512³ voxels)	0.300
lmd_ssl_v2_neuron_instance (whole volumes)	0.277

Per-volume scores and the other models are on the leaderboard.

License

DINOv3 License (LICENSE.md). This model is a derivative of Meta's DINOv3 weights and code, so that agreement governs its use and redistribution. 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.
• A. Hatamizadeh et al. UNETR: Transformers for 3D Medical Image Segmentation. WACV 2022.
• A. Sheridan et al. Local shape descriptors for neuron segmentation. Nature Methods 20, 2023.
• S. Wolf et al. The Mutex Watershed and its Objective: Efficient, Parameter-Free Graph Partitioning. TPAMI 2020.