Spatial Masked-Set Learning for Sparse Multi-Shell Diffusion MRI Signal Synthesis

This repository contains the official implementation of spatial masked-set SHORE prediction. Our paper, "Spatial Masked-Set Learning for Sparse Multi-Shell Diffusion MRI Signal Synthesis," addresses the long acquisition times of dense multi-shell diffusion MRI by recovering a continuous q-space representation from a small, arbitrary subset of measured gradients.

Yousef Sadegheih, Pratibha Kumari, and Dorit Merhof

---

📑 Table of Contents

1. Abstract
2. Updates
3. Key Contributions
4. Model Architecture
5. Dataset and Pre-trained Weights
6. Results
7. Getting Started
   • Requirements
   • Installation
   • Data Preparation
   • Training
   • Evaluation
   • Notes
8. Repository Structure
9. Acknowledgments
10. Citation

---

📝 Abstract

Dense multi-shell diffusion MRI provides rich q-space information but requires long acquisition times. We propose a spatial masked-set framework for sparse multi-shell diffusion MRI signal synthesis. The model treats observed measurements as an unordered set, uses a local 3 × 3 × 3 neighborhood for spatial context, and predicts radial-order-6 SHORE coefficients for the center voxel. The coefficients can then be decoded analytically to synthesize signals at arbitrary q-space locations. Training combines shell-wise gradient dropping, dense signal supervision, and rotation-consistent SHORE targets so that sparse input signals remain aligned with their coefficient supervision under augmentation. We evaluate on held-out HCP100 white-matter voxels by retaining limited subsets of measured diffusion-weighted signals from the reference acquisition. The proposed method achieves lower signal NMSE than both analytical q-space models and a state-of-the-art continuous dMRI signal synthesis model designed for arbitrary input and output q-space sampling. In the b = 1000 setting with 10 input gradients, it achieves 2.70% NMSE, a 22.4% relative reduction over this continuous model. Fractional anisotropy on reconstructed b = 1000 signals provides a complementary tensor-derived endpoint, with analytical models remaining competitive for FA despite higher dense-signal NMSE across the evaluated q-space.

---

🔔 Updates

• 😎 First release — code, pre-trained weights, and preprocessed dataset – August 7, 2026
• 🥳: Paper accepted at MICCAI Workshop 2026 – August 6, 2026

---

⚡ Key Contributions

• Masked-set SHORE prediction for arbitrary sparse multi-shell inputs. Observed measurements are treated as an unordered set with explicit q-space coordinates, so a single trained model handles any input gradient subset without a fixed number or ordering of measurements.
• A compact spatial architecture using 3 × 3 × 3 context before q-space pooling. Local anatomical coherence is injected into the center-voxel features through mask-aware spatial attention, before permutation-invariant pooling over the observed q-space samples.
• Sparse-input, dense-output training with rotation-consistent SHORE supervision. Shell-wise gradient dropping simulates reduced protocols while supervision stays dense, and SHORE targets are re-solved after each random gradient rotation so the augmented input and target share one acquisition geometry.
• Evaluation on HCP100 showing lower signal NMSE than a state-of-the-art continuous dMRI reconstruction model and analytical q-space models across the evaluated sparse single- and multi-shell signal-generation settings.

---

⚙️ Model Architecture

Each measurement is described by its normalized q-space position, b-value, and signal intensity, then encoded by a shared MLP. Center-voxel features are refined with mask-aware attention over the 26 neighbors measured at the same encoding, pooled over the observed q-space set with mean–max pooling, and mapped to 50 standardized SHORE coefficients. The coefficients are denormalized and decoded with the SHORE basis to synthesize dense diffusion signals at arbitrary query locations.

The final network contains approximately 0.23 M trainable parameters, comparable to the 0.25 M parameters of DISCUS.

For a detailed explanation of each component, please refer to our paper.

---

🗄️ Dataset and Pre-trained Weights

Experiments use the HCP100 subset of the Human Connectome Project. Each subject has three nonzero shells at approximately b = 1000, 2000, and 3000 s/mm², with 90 diffusion-weighted measurements per shell and 18 b = 0 volumes. Subjects are split at the subject level into 80 training, 10 validation, and 10 test subjects; the exact split is checked in as whole_subject_list.json.

Resource	Download
Pre-trained weights	[Download Weights]
Preprocessed HCP100 (compact format)	[Download Dataset]

The released checkpoint is the model reported in the paper: 229,651 trainable parameters, selected at 250,000 steps by subject-macro validation signal MSE (0.003667). It is a full training checkpoint, so it also carries the optimizer and scheduler state and can be used to resume training via --load_checkpoint_path. Its configuration is checked in at configs/pretrained_experiment.json — the experiment.json of the original run, with machine-specific paths blanked — so the weights can be evaluated without training anything first. See Evaluation.

The preprocessed dataset is the output of Data Preparation for all 100 subjects, so downloading it lets you skip preprocessing entirely. It is distributed as a multi-part zip archive; all parts must be placed in the same directory before extracting:

# all four parts must be present: .z01, .z02, .z03 and .zip
zip -s0 HPC100_whole_dataset_compact_3d_only.zip --out joined.zip
unzip joined.zip

7-Zip and WinRAR open the split set directly, without the rejoin step. Extraction yields one directory per subject; point data_path in train.sh at the extracted HPC100_whole_dataset_compact_3d_only/.

Notes

• HCP data requires a data use agreement and must be obtained from the HCP data portal. The released archive contains derived quantities (b0-normalized signals, SHORE coefficients, and gradient tables) rather than the raw imaging data.
• White-matter masks are obtained with FSL BET and FAST via preprocess/make_wm_masks.sh; the preprocessing script expects them at <subject>/T1w/WM_mask_in_DWI.nii.gz.
• The released dataset contains only what training and evaluation read, so it works with --dataset_type base_final_ram_compact_neighbor (the default) and base_final_mmap_compact_neighbor. The non-compact base_final_mmap_neighbor option additionally needs per-subject dense patch arrays, which are not distributed.

---

📈 Results

Dense multi-shell reconstruction from sparse b = 1000 inputs

Evaluated on the HCP100 test set. The gradient count denotes the number of observed b = 1000 directions. Signal NMSE is computed on the reconstructed b = 1000, b = 2000, and b = 3000 signals with 90 directions per shell, while MSE_FA is computed from FA maps fitted to the reconstructed b = 1000 shell. Both metrics are reported as mean (standard deviation) in percentages. Bold marks the best and italic the second-best mean for each sampling budget.

# input directions	DISCUS MSEFA	DISCUS NMSE	SHORE MSEFA	SHORE NMSE	MAP MSEFA	MAP NMSE	Ours MSEFA	Ours NMSE
5	2.15 (0.043)	4.81 (0.011)	1.56 (0.081)	12.70 (0.768)	3.19 (0.153)	7.29 (0.214)	1.80 (0.042)	4.08 (0.004)
10	0.83 (0.038)	3.48 (0.022)	0.66 (0.069)	10.80 (0.701)	1.05 (0.051)	5.34 (0.215)	0.41 (0.028)	2.70 (0.003)
15	0.61 (0.076)	3.18 (0.035)	0.41 (0.046)	10.38 (0.647)	0.54 (0.036)	4.89 (0.221)	0.27 (0.013)	2.49 (0.002)
20	0.50 (0.087)	3.05 (0.037)	0.29 (0.035)	10.11 (0.663)	0.34 (0.033)	4.68 (0.199)	0.22 (0.010)	2.40 (0.003)
25	0.43 (0.083)	2.97 (0.045)	0.22 (0.028)	9.97 (0.655)	0.25 (0.028)	4.57 (0.197)	0.17 (0.010)	2.34 (0.003)
30	0.39 (0.076)	2.94 (0.040)	0.17 (0.021)	9.82 (0.630)	0.21 (0.022)	4.51 (0.186)	0.17 (0.011)	2.31 (0.003)
40	0.36 (0.063)	2.92 (0.040)	0.11 (0.013)	9.65 (0.603)	0.18 (0.015)	4.41 (0.182)	0.15 (0.015)	2.27 (0.004)
50	0.36 (0.043)	2.93 (0.039)	0.07 (0.008)	9.54 (0.600)	0.20 (0.011)	4.36 (0.174)	0.14 (0.015)	2.24 (0.005)

Signal NMSE across sparse single- and multi-shell input settings

Each panel specifies the shells available as sparse input; an x-axis value of k means that k directions are retained from each included shell using antipodal sampling. Lower is better.

Ablation study

NMSE in percent, lower is better. Full observation uses all 270 diffusion-weighted measurements, while sparse input uses b = 1000 with 10 retained gradients.

Setting	Variant	NMSE (%) ↓
Full observation	Fixed-order MLP	2.25
Full observation	MLP with shuffled gradient order	7.32
Full observation	Basic Deep Set	2.15
Full observation	Basic Deep Set + rotation augmentation	2.07
Full observation	Spatial masked-set model + rotation augmentation	1.90
Sparse input	Without gradient dropping during training	6.65
Sparse input	With gradient dropping during training	2.70

This repository ships the final model only. The ablation variants above are reported for reference and are not included in the release.

---

🚀 Getting Started

🛠️ Requirements

• Operating System: Ubuntu 22.04 or higher
• Python: 3.10
• CUDA: Version 12.x
• Package Manager: Conda
• FSL: required for white-matter mask generation only, not for training or evaluation
• Hardware:
  • GPU with 16 GB memory or larger (recommended)
  • For our experiments, we used a single NVIDIA A100 with 40 GB VRAM; the final model takes approximately 10 GPU-hours
  • The RAM-resident dataset preloads the training split; use --dataset_type base_final_mmap_compact_neighbor on hosts with limited system memory

📦 Installation

To create the environment and install all dependencies, run:

./env_creation.sh

This creates a Conda environment named shorepred, installs a CUDA build of PyTorch, and installs the remaining requirements. Adjust the PyTorch index URL inside the script to match your CUDA version. Alternatively, install into an existing environment with pip install -r requirements.txt.

📂 Data Preparation

Download the HCP100 subjects and arrange them in the standard HCP layout:

<raw_root>/<subject_id>/T1w/T1w_acpc_dc_restore_1.25.nii.gz
<raw_root>/<subject_id>/T1w/Diffusion/data.nii.gz
<raw_root>/<subject_id>/T1w/Diffusion/bvals
<raw_root>/<subject_id>/T1w/Diffusion/bvecs

Step 1 — white-matter masks. Training and evaluation are restricted to white matter, so each subject needs a WM_mask_in_DWI.nii.gz. This requires FSL on your PATH:

./preprocess/make_wm_masks.sh /path/to/HCP100_raw --skip-existing

The script runs FSL bet for brain extraction and fast for three-class tissue segmentation on the HCP 1.25 mm T1, thresholds the white-matter partial-volume estimate at 0.5, and restricts the result to the brain mask. The 1.25 mm T1 is already aligned with the diffusion data, so no registration or resampling into diffusion space is required. Each subject gains:

<raw_root>/<subject_id>/T1w/WM_mask_in_DWI.nii.gz

Step 2 — compact dataset. Then run preprocessing:

./preprocess.sh

or equivalently:

python preprocess/preprocess_raw_to_compact.py \
    --raw_root  /path/to/HCP100_raw \
    --output_root /path/to/HCP100_compact \
    --skip_existing

For each subject this normalizes the diffusion signal voxel-wise by the mean b = 0 signal, extracts the 3 × 3 × 3 white-matter neighborhoods, fits radial-order-6 SHORE coefficients to every center voxel as the dense supervision target, precomputes the SHORE basis matrix, and generates the antipodal sparse-sampling bank used to simulate reduced protocols. Signal voxels shared between overlapping patches are stored once, with patch-index maps recovering the original neighborhoods — on HCP100 white matter this stores roughly a tenth of the voxel slots a dense patch array would need.

Each output subject directory contains:

<subject_id>_meta.npz                  bvals, bvecs, and the b0 volume
<subject_id>_shore_matrix.npy          SHORE basis matrix, (G, 50)
<subject_id>_signals_unique.npy        deduplicated normalized signals
<subject_id>_signal_coords.npy         voxel coordinate per unique signal
<subject_id>_patch_indices_3x3x3.npy   3x3x3 patch -> unique signal index, (V, 27)
<subject_id>_coords_3x3x3.npy          valid 3x3x3 patch centers, (V, 3)
<subject_id>_shore_coeffs_3x3x3.npy    center-voxel SHORE targets, (V, 50)
<subject_id>_antipodal_indices.pt      sparse-sampling index bank
<subject_id>_compact_manifest.json     provenance and verification record

The SHORE normalization statistics used to standardize the coefficient targets are checked in under normalization_stats/ and are selected automatically from the subject list. To recompute them for a different cohort, use:

python analysis/compute_cohort_shore_norm.py --help

🏋️ Training

Update data_path and checkpoint_path in train.sh, then run:

./train.sh

The script reproduces the paper configuration: the spatial masked-set network trained for 250,000 steps with Adam, learning rate 10⁻³, weight decay 10⁻⁴, a polynomial learning-rate schedule with power 0.9, batch size 512, loss weights λ_c = 1 and λ_s = 10, shell-wise gradient dropping, and random rotation augmentation applied with probability 0.25.

Any flag can be overridden on the command line, since arguments are forwarded to main.py:

./train.sh --train_batch_size 256 --n_steps 90000 --name my_run

Each run writes to <checkpoint_path>/<network>-<trainer>-wb/<experiment_name>/, where <experiment_name> encodes the configuration. For the paper settings this is:

checkpoints/spatialAttLiteStrong-mask-wb/
└── lr0.001_bs512_st250k-decay_cmb-shw1-sgw10-asm_dswhole-compact-ram-drop-rot0.25_bval_ctx1/
    ├── experiment.json      complete argument set, parameter count, FLOPs
    ├── checkpoints/
    │   ├── checkpoint_epoch_<epoch>.pth
    │   └── best_checkpoint_epoch_<epoch>_MSE_<val_metric>.pth
    ├── results/
    │   ├── test_results_epoch_<epoch>.json
    │   └── test_results_best_model_epoch_<epoch>_MSE_val_<val_metric>.json
    ├── logs/train.log
    └── tensorboard/

experiment.json is written at the top of the run, next to the checkpoints/ directory. It records every argument the run was launched with, plus the trainable parameter count and multiply-adds, and it is what evaluation/evaluate_sparse.py reads to rebuild the model. The copy for the released weights is checked in at configs/pretrained_experiment.json.

Checkpoint selection uses subject-macro signal MSE on validation with 15 gradients from all three shells. After training, the final checkpoint is tested, then the best checkpoint is loaded and tested again.

📊 Evaluation

To reproduce the sparse evaluation grid, point evaluate.sh at a trained checkpoint and the experiment.json sitting beside it in the run directory:

./evaluate.sh

or equivalently:

python evaluation/evaluate_sparse.py \
    --checkpoint /path/to/best_model.pth \
    --experiment_config /path/to/experiment.json \
    --data_path /path/to/HCP100_compact \
    --output results/sparse_test.json \
    --device cuda --batch_size 512 --amp 1

This evaluates every non-empty combination of the b = 1000, 2000, and 3000 shells at each retained-gradient count k ∈ {5, 10, 15, 20, 25, 30, 40, 50}, and writes a JSON file keyed by <k>_<shells>. The architecture is rebuilt from experiment.json, so no architecture flags need to be repeated.

To evaluate the released pre-trained weights instead of your own run, point --experiment_config at the checked-in configuration:

python evaluation/evaluate_sparse.py \
    --checkpoint /path/to/best_checkpoint_epoch_25_MSE_0.003667.pth \
    --experiment_config configs/pretrained_experiment.json \
    --data_path /path/to/HCP100_compact \
    --output results/sparse_test.json \
    --device cuda --batch_size 256 --amp 1

Reported signal NMSE is available in each entry as NMSE_signal_patient_macro. For example, results["10_1000"]["NMSE_signal_patient_macro"] reproduces the 2.70 % entry in the table above.

⚠️ Notes

• Path configuration: update the paths at the top of preprocess.sh, train.sh, and evaluate.sh before running them.
• Normalization contract: every experiment intentionally loads the rotation-aware SHORE normalization archive, including runs with --rotate_bvecs 0, so that training, validation, and testing share one normalization contract. Actual b-vector rotation still occurs only for training samples.
• Outlier handling: final evaluation logs voxels whose per-voxel signal MSE exceeds 10 and excludes them from metric aggregation, as a safeguard against known preprocessing outliers.
• Comparison methods: DISCUS was trained from its public source using the hyperparameters provided by its authors, and is not redistributed here. Analytical SHORE and MAP-MRI baselines were evaluated with DIPY 1.9.0.
• Antipodal sampling bank: the sparse masks are drawn from a precomputed bank of candidate subsets so the same mask distribution is shared across methods and runs. Generation is deterministic given the gradient table — one greedy repulsion subset per seed direction, deduplicated — so it introduces no run-to-run variability. Regenerate it for an existing dataset with python preprocess/antipodal.py --data_root /path/to/HCP100_compact --overwrite.

---

📁 Repository Structure

├── main.py                          training entry point
├── utils.py                         argument parsing, dataloaders, input assembly
├── train.sh / evaluate.sh / preprocess.sh
├── models/
│   └── deep_set_lite_strong.py      spatial masked-set SHORE predictor
├── trainers/
│   ├── base_trainer_final.py        training loop, checkpointing, logging
│   └── base_trainer_final_mask.py   sparse-input, dense-output training
├── dataset/
│   ├── dataset_mmap.py              base dataset, masking and rotation
│   ├── dataset_mmap_compact.py      compact storage, memory-mapped
│   └── dataset_ram_compact.py       compact storage, RAM-resident
├── losses/
│   ├── losses_final.py              coefficient and signal losses
│   └── shore_matrix_pytorch.py      differentiable SHORE basis
├── evaluation/
│   ├── eval_final.py                NMSE, FA, and patient-macro metrics
│   └── evaluate_sparse.py           full sparse shell/count evaluation grid
├── preprocess/
│   ├── make_wm_masks.sh             FSL white-matter mask generation
│   ├── preprocess_raw_to_compact.py raw HCP NIfTI -> compact training format
│   ├── antipodal.py                 antipodal repulsion sampling banks
│   └── common.py                    shared constants and helpers
├── analysis/
│   └── compute_cohort_shore_norm.py rotation-aware SHORE normalization stats
└── normalization_stats/             checked-in SHORE normalization archives

---

🤝 Acknowledgments

This work builds on DIPY for SHORE, MAP-MRI, and tensor fitting, and compares against DISCUS by Ewert, Kügler, Stirnberg, Koch, Yendiki and Reuter. Data were provided by the Human Connectome Project, WU-Minn Consortium. We thank the authors of these projects for making their work available.

---

📚 Citation

If you find this work useful for your research, please cite:

@inproceedings{sadegheih2026spatial,
  title={Spatial Masked-Set Learning for Sparse Multi-Shell Diffusion MRI Signal Synthesis},
  author={Sadegheih, Yousef and Kumari, Pratibha and Merhof, Dorit},
  year={2026},
  organization={Springer}
}