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