BrainAxL

Official weights for BrainAxL, a multiscale axial-LSTM foundation model for 3D brain MRI.

• Source code
• Training guide
• Sanitized training logs
• Citation metadata

This repository contains only the canonical self-supervised BrainAxL-B checkpoint from run 19726. FOMO26 downstream and submission-specific weights are intentionally not part of this foundation-model release.

Model

BrainAxL is a five-stage 3D U-Net encoder with channels [40, 80, 160, 320, 640]. Bidirectional LSTMs scan all three spatial axes at the two deepest stages. Pretraining combines masked 3D reconstruction with variance/covariance regularization of a 512-dimensional semantic projection.

Field	Value
Input	single-channel 64 x 64 x 64 training crop
Encoder parameters	31,837,640
Active pretraining parameters	46,308,601
Mask ratio	60% in 4 x 4 x 4 blocks
Optimizer	AdamW, lr 1e-4, weight decay 3e-5
Schedule	cosine, no warm-up
Training	20 epochs / 34,660 updates
Precision	bfloat16 mixed

Files

Path	Purpose
brainaxl-b/model.safetensors	Canonical safe tensor-only state dictionary
brainaxl-b/model.ckpt	Weight-only Lightning wrapper for Asparagus
brainaxl-b/training_config.yaml	Recovered Hydra training configuration
brainaxl-b/hparams.yaml	Lightning hyperparameters
brainaxl-b/hydra_overrides.yaml	Exact reported command overrides
config.json	Architecture metadata
preprocessing.json	Pretraining input contract
metadata.json	Provenance, sizes, and SHA-256 hashes
checksums.sha256	Download integrity checks

Load

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

path = hf_hub_download(
    repo_id="hieuphamha/BrainAxL",
    filename="brainaxl-b/model.safetensors",
)
state_dict = load_file(path, device="cpu")

Install the linked GitHub package and instantiate brainaxl_b to load this state dictionary. Existing Asparagus pipelines can instead download brainaxl-b/model.ckpt and pass it as checkpoint_path.

Training data and evaluation

The checkpoint was pretrained on the FOMO300K corpus under the FOMO26 Methods Track rules. No downstream labels were used during pretraining. The public GitHub repository provides the exact objective, configuration, training commands, a 20-epoch sanitized log, and downstream evaluation methodology.

Intended use

Appropriate uses include research on 3D brain MRI representation learning and initialization of validated downstream models. This model is not an approved medical device and must not be used for clinical diagnosis or treatment without task-specific validation and appropriate regulatory review.

Limitations

Performance can shift with anatomy, sequence, acquisition protocol, scanner, population, preprocessing, and label policy. The checkpoint uses historical state-dictionary names for compatibility with the reported run. Bitwise reproduction can depend on CUDA, PyTorch, data ordering, and preprocessing versions.

License and citation

Code is Apache-2.0. Weight use is subject to the applicable FOMO26 and FOMO300K terms; see WEIGHTS_LICENSE.md. Citation metadata is maintained in the GitHub repository.