BrainAxL
Official weights for BrainAxL, a multiscale axial-LSTM foundation model for 3D brain MRI.
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.