---
library_name: pytorch
tags:
- medical-imaging
- mri
- self-supervised-learning
- foundation-model
- 3d
- fomo26
license: other
datasets:
- FOMO-MRI/FOMO300K
pipeline_tag: image-feature-extraction
---

# BrainAxL

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

- [Source code](https://github.com/hieuphamha19/BrainAxL)
- [Training guide](https://github.com/hieuphamha19/BrainAxL/blob/main/docs/training.md)
- [Sanitized training logs](https://github.com/hieuphamha19/BrainAxL/tree/main/training_logs)
- [Citation metadata](https://github.com/hieuphamha19/BrainAxL/blob/main/CITATION.cff)

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

~~~python
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.