Add 3D volume modality

[HarshaSatyavardhan]

Summary

This PR adds native 3D medical volume support for Stage-1 alignment experiments.

• Added 3D volume dataset loaders:
  • fs-volume for loading .npy volumes from filesystem paths.
  • raw-volume for loading in-memory / Arrow-style volume payloads.
• Added VolumeProcessor for 3D medical volume preprocessing:
  • accepts .npy volume arrays.
  • normalizes to [0, 1].
  • resizes to (1, 32, 256, 256) contract.
• Added VolumeConfig:
  • default encoder: GoodBaiBai88/M3D-CLIP.
  • volume_size=(32,256,256).
  • patch_size=(4,16,16).
  • pool_factor=(2,2,2).
  • validates derived token path: 2048 -> 256.
• Added volume_3d modality:
  • loads a pretrained 3D vision encoder.
  • supports encode_image() or vision_model(...).last_hidden_state.
  • uses M3D-style 3D spatial pooling from 2048 patch tokens to 256 visual tokens.
  • keeps the MultiMeditron MLPProjector for mapping encoder features to LLM hidden size.
• Added Stage-1 alignment cookbook config:
  • freezes the LLM.
  • freezes the 3D vision encoder.
  • trains only the projector.
• Updated dataset checking support for 3D image modality inputs.

Validation

• Config/grid smoke test:
  • (32,256,256) with patch (4,16,16) produces (8,16,16) -> (4,8,8).
  • proj_out_num=256.
• Spatial pooling smoke test:
  • input tokens: (2,2048,768).
  • output tokens: (2,256,768).
• Real M3D-CLIP forward smoke test:
  • output shape: (1,256,4096).
• 3D MedMNIST validation completed (on NoduleMNIST3D)
• Training/eval completed successfully with:
  • TinyLlama/TinyLlama-1.1B-Chat-v1.0
  • EPFLiGHT/Meditron3-Gemma2-2B

Training with a Dataset Status

The original M3D alignment dataset, M3D-Cap, currently shows disabled access due to a DMCA/licensing discussion: M3D-Cap discussion #4.

For follow-up alignment experiments, CT-RATE is a practical open alternative because it provides 3D chest CT volumes paired with radiology text reports. CT-RATE is suitable for projector alignment experiments where the LLM and vision encoder are frozen and only the projector is trained.

Spatial pooling

Input volume: 32 × 256 × 256
Patch size:   4 × 16 × 16

Patch grid:
32/4   = 8
256/16 = 16
256/16 = 16

So total patch tokens:
8 × 16 × 16 = 2048 tokens

after the 3D encoder:

2048 visual tokens, each token has 768 features
Shape: 2048 × 768

Then 3D spatial pooling reduces the spatial grid:

8 × 16 × 16  →  4 × 8 × 8

So token count becomes:

4 × 8 × 8 = 256 tokens

After spatial pooling:

256 visual tokens, each token still has 768 features
Shape: 256 × 768

M3D-CLIP outputs:

hidden size = 768

But Llama-3.1-8B uses:

hidden size = 4096

So the MLP projector converts each visual token from M3D-CLIP space into Llama embedding space:

768 → 4096

References

• M3D paper: M3D: Advancing 3D Medical Image Analysis with Multi-Modal Large Language Models
• M3D-Cap dataset: GoodBaiBai88/M3D-Cap
• M3D-Cap licensing discussion: Discussion #4
• CT-RATE dataset: ibrahimhamamci/CT-RATE
• CT-RATE paper: A foundation model utilizing chest CT volumes and radiology reports for supervised-level zero-shot detection of abnormalities
• 3D MedMINIST dataset

[HarshaSatyavardhan]

Hi @MichelDucartier, I have implemented the changes you suggested and updated this PR accordingly.

Could you please take a look when you get a chance and let me know if this direction looks good, or if any further changes are needed?