Herbert23H/V2VA_Pretrain_Checkpoint

V2VA — Volume Token Volume Alignment (ViT-B/16, MedMNIST+3D 64³)

Volume Token Volume Alignment across Embedding and Volumetric Spaces for 3D Medical Image Classification — accepted at PRICAI 2026 (Pacific Rim International Conference on Artificial Intelligence).

Zhenghao Duan, Dongyun Lin, Weiyao Lan — School of Aerospace Engineering, Xiamen University, Xiamen, China.

📄 The official proceedings have not been published yet; this model card will link the paper and DOI once they are available.

Self-supervised 3D medical volume encoder, produced by the V2VA project, built on top of facebookresearch/jepa (V-JEPA).

This is the pretrained encoder only. It is used as a frozen backbone for downstream MedMNIST+3D classification, where the V2VA head re-folds the encoder tokens [B, N, E] back into a voxel grid and classifies with a residual 3D CNN.

Files

File	Size	SHA-256
jepa-latest.pth.tar	1.5 GB	a731d61d67faf06c094b3ba002054103ce5c92c45568577a255030ff16307a5b

The archive is a PyTorch checkpoint dict with the keys:

epoch, encoder, predictor, target_encoder, opt, scaler

target_encoder holds the EMA target encoder — this is the representation used by the downstream evaluations (checkpoint_key: target_encoder in the V2VA eval configs). opt / scaler are only needed to resume pretraining.

import torch

ckpt = torch.load("jepa-latest.pth.tar", map_location="cpu", weights_only=False)
print(ckpt["epoch"])
target_sd = ckpt["target_encoder"]   # frozen encoder used downstream

The codebase loads this file directly — see evals/mri_classification_frozen/eval.py.

How it was trained

Self-supervised (no labels) V-JEPA-style masked latent prediction on 64³ CT volumes from MedMNIST+3D — organmnist3d_64, nodulemnist3d_64, fracturemnist3d_64.

Setting	Value
Backbone	ViT-B/16 (vit_base), 3D tubelet
Input	64³ single-channel volumes, num_slices=64
Patch / tubelet	patch_size=4, depth_patch_size=4
Predictor	depth 6, embed dim 384, use_mask_tokens=true, zero_init_mask_tokens=true
Masking	multiblock3d_mri, 8 blocks, two scale sets (depth-heavy 0.75 / space-heavy 0.75)
Epochs	300
LR	2e-4 → 1e-5, warmup 10, weight decay 0.02 → 0.2, cosine schedule
EMA	0.998 → 1.0
Batch size	8
Precision	bfloat16
Seed	234

Full resolved config: experiments/params-pretrain.yaml.

Downstream results

Frozen encoder + a lightweight head, averaged over the six MedMNIST+3D tasks:

Head	Params (M)	FLOPs (G)	Inference (ms)	Avg. AUC (%)
Average pooling	23.63	11151.9	532.1	84.1
Attentive pooling	38.11	11307.4	535.6	84.5
CLS token	37.81	13009.8	611.6	84.2
Token2Volume (V2VA)	27.14	11673.4	541.1	92.2

Params / FLOPs / latency measured on adrenalmnist3d_64. Figures: trade-off plot, per-task radar.

Intended use and limitations

• Intended for research on 3D medical image representation learning.
• Pretrained on MedMNIST+3D, which contains small, downsampled 64³ CT volumes. It is not a clinical-grade model and must not be used for diagnosis.
• Evaluate on your own data before any downstream use; the encoder was trained without labels and carries no disease-specific supervision.

License

CC BY-NC 4.0 (Attribution-NonCommercial 4.0 International) — see LICENSE.

These weights are a derivative work of facebookresearch/jepa, Copyright © Meta Platforms, Inc. and affiliates, which is licensed under CC BY-NC 4.0. Non-commercial use only. MedMNIST+ is distributed under CC BY 4.0; please cite it when using this model.

Citation

If you use these weights, please cite the paper (accepted at PRICAI 2026; the official proceedings are not published yet, so page numbers and DOI are still unknown):

@inproceedings{duan2026v2va,
  title     = {V2VA: Volume Token Volume Alignment across Embedding and Volumetric Spaces for 3D Medical Image Classification},
  author    = {Duan, Zhenghao and Lin, Dongyun and Lan, Weiyao},
  booktitle = {Pacific Rim International Conference on Artificial Intelligence (PRICAI)},
  year      = {2026},
  note      = {Accepted; official proceedings not yet published}
}