DOCSEARCHbyAlgolia
Submit Docs

4M: Massively Multimodal Masked Modeling

repository·main·Indexed 23 days ago

https://github.com/apple/ml-4m

A framework for training any-to-any multimodal foundation models using tokenization and masking. 4M enables flexible and steerable multimodal generation across diverse modalities and tasks. It includes tools for loading foundation models (B, L, and XL variants) and specialist models for text-to-image and super-resolution from the Hugging Face Hub or local safetensors, as well as a Demo4MSampler wrapper for RGB-to-all and caption-to-all generation workflows.

Tokens
9.3K
Snippets
27
Records
41
Agent score
83%

What's inside 4M (fourm)

  1. Understand the 4M training repository structure

    main

    The 4M training codebase is organized into two primary functional areas:

    fourm/models/ (Model Architecture)

    Contains the core logic for 4M models:

    • models/fm.py: Defines the FourM module, including architecture and forward pass logic.
    • models/encoder_embeddings.py & models/decoder_embeddings.py: Handle per-modality mapping of tokens/patches to embeddings and embeddings to logits, including positional and modality embeddings.
    • fm_vit.py: Provides the FourMViT module for RGB-only ViT behavior.
    • generate.py: Contains sampling logic and utilities for any-to-any generation.

    fourm/data/ (Data Pipeline)

    Handles the multimodal data lifecycle:

    • data/modality_info.py: Defines modality metadata (name, type, vocabulary size, etc.).
    • data/unified_datasets.py: Loads aligned multimodal datasets from local paths or cloud stores like S3.
    • data/modality_transforms.py: Manages aligned data augmentations via UnifiedDataTransform and per-modality preprocessing.
    • data/masking.py: Implements multimodal input/target masking using token budgets and Dirichlet sampling.

  2. Understand the 4M Tokenization Structure

    main

    Tokenization in 4M converts diverse modalities (images, feature maps, sequences) into a unified representation space. The core logic resides in the fourm/vq/ directory:

    • fourm/vq/models/: Encoder and decoder architectures.
    • fourm/vq/percept_losses/: Perceptual loss implementations for VQ-VAE training.
    • fourm/vq/quantizers/: Quantizer implementations.
    • fourm/vq/scheduling/: Diffusion schedules and inference pipelines.
    • fourm/vq/vqvae.py: Definitions for standard VQ-VAE and diffusion-based VQ-VAE classes.
    • fourm/vq/__init__.py: Contains get_image_tokenizer for autoloading tokenizers.

    Training scripts are located in the root directory:

    • run_training_vqvae.py: For standard VQ-VAEs.
    • run_training_divae.py: For diffusion-based VQ-VAEs.
    • run_training_vqcontrolnet.py: For ControlNet detokenizers.

  3. Configure chained generation schedules

    main

    Chained generation involves generating modalities one-by-one, conditioning each subsequent modality on the previously generated ones. This is managed via utils/generation.py using a generation schedule and a Sampler.

    Input Format

    Inputs are dictionaries of modalities. Each modality dictionary contains:

    • Tokens: A set of tokens containing data or placeholder values.
    • Input mask: Specifies which parts of the tokens are used as input.
    • Target mask: Specifies which parts of the tokens are to be predicted.

    Domain Configuration

    • cond_domains: Domains used for conditioning.
    • target_domains: Domains to be predicted. The order determines the sequence of prediction.

    Hyperparameter Granularity

    Most parameters can be applied to all target modalities at once, or specified per modality using hyphens in configs or lists in the build_chained_generation_schedules function.

    Example: For a caption $\rightarrow$ CLIP $\rightarrow$ RGB chain, setting temps: 3.0-0.5 applies a temperature of 3.0 to CLIP and 0.5 to RGB.

  4. Install 4M (fourm)

    main

    To install the 4M framework, clone the repository, create a new conda environment with Python 3.9, and install the package in editable mode. It is recommended to upgrade pip first to ensure PEP 660 support.

    git clone https://github.com/apple/ml-4m
cd ml-4m
conda create -n fourm python=3.9 -y
conda activate fourm
pip install --upgrade pip
pip install -e .

  5. Perform chained multimodal generation with run_generation.py

    main

    Use the run_generation.py script to automate chained X→Y→etc... generation on a dataset. The script requires configuration files for the model, data, base generation settings, and optional super-resolution settings.

    To run generation on 8 GPUs for a text→CLIP→RGB pipeline using 4M-XL on Parti prompts with super resolution, use the following command structure:

    OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_generation.py \
  -c cfgs/default/generation/models/4m-xl_mod7+sr_4m-l_mod7.yaml \
  -dc cfgs/default/generation/data/parti_3x.yaml \
  -gc cfgs/default/generation/settings_base/T2CR_roar49-25_cfg3_t6-0.5.yaml \
  -src cfgs/default/generation/settings_sr/x2CR_mg8_cfg3_t1const.yaml

    Note: Ensure 4M and tokenizer checkpoints are either downloaded and correctly pointed to in the configs, or available via Hugging Face Hub.

    OMP_NUM_THREADS=1 torchrun --nproc_per_node=8 run_generation.py -c cfgs/default/generation/models/4m-xl_mod7+sr_4m-l_mod7.yaml -dc cfgs/default/generation/data/parti_3x.yaml -gc cfgs/default/generation/settings_base/T2CR_roar49-25_cfg3_t6-0.5.yaml -src cfgs/default/generation/settings_sr/x2CR_mg8_cfg3_t1const.yaml

  6. Organize large datasets using the Modified WebDataset format

    main

    For large-scale training or datasets stored in cloud object stores (like S3), use a modified WebDataset format. This involves splitting data into tarfiles containing 1,000 to 10,000 samples per modality. This structure reduces object read requests when streaming from the cloud.

    Key Requirements for Alignment:

    • Identical Filenames: Files across different modality folders must have the same base name. Only the modality folder name and the file extension should differ.
    • Numerical Ordering: Shards must be ordered numerically (e.g., shard-00000.tar, shard-00001.tar) to support brace-expand notation.
    • Modality Folders: Each modality (e.g., rgb, caption, depth) should have its own root-level directory containing its respective tarfiles.

    Example Structure:

    root/rgb/shard-00000.tar
root/rgb/shard-00001.tar

root/caption/shard-00000.tar
root/caption/shard-00001.tar
    root/modality_a/shard-00000.tar
root/modality_a/shard-00001.tar

root/modality_b/shard-00000.tar
root/modality_b/shard-00001.tar

  7. Train a Diffusion Decoder with a Frozen Encoder

    main

    When training a diffusion decoder on top of a pre-trained VQ-VAE encoder, it is recommended to train the standard VQ-VAE first to establish a stable representation. To freeze the encoder during the subsequent diffusion decoder training, include these keys in your configuration:

    • full_ckpt: Path to the pre-trained checkpoint (.pth).
    • freeze_enc: Set to True to prevent encoder updates.
    • input_size_enc: The size of the encoder positional embeddings.
    full_ckpt: /path/to/checkpoint.pth
freeze_enc: True # Decoder can be trained from scratch or fine-tuned without the encoder
input_size_enc: 256 # Size of the encoder positional embeddings

  9. Expose 4M environment to Jupyter Notebooks

    main

    If you want to use the fourm conda environment as a kernel in Jupyter, install ipykernel and register the kernel:

    pip install ipykernel
python -m ipykernel install --user --name fourm --display-name "4M (fourm)"

  10. Organize small local datasets using the Simple Hierarchical format

    main

    For smaller datasets, such as validation or transfer sets stored locally, use a simple hierarchical structure. Alignment is maintained by ensuring that folder and file names are identical across all modality directories.

    Example Structure:

    root/modality_a/folder_x/xxx.ext
root/modality_a/folder_y/xxy.ext

root/modality_b/folder_x/xxx.ext
root/modality_b/folder_y/xxy.ext
    root/modality_a/folder_x/xxx.ext
root/modality_a/folder_y/xxy.ext
root/modality_a/folder_z/xxz.ext

root/modality_b/folder_x/xxx.ext
root/modality_b/folder_y/xxy.ext
root/modality_b/folder_z/xxz.ext

  11. Best practices for 4M generation

    main

    Modality Transitions

    • Intermediate Steps: When generating dense modalities (like RGB) from abstract ones (like text), use intermediate steps (e.g., Text $\rightarrow$ CLIP $\rightarrow$ RGB) to improve quality.
    • Image Schemes: Use maskgit for simpler, more controllable images; use roar (Random Order Auto Regression) for more diverse images.

    Sampling Hyperparameters

    • Temperature: For the first few tokens, use a higher temperature that decays over the schedule. Early stages decide most rough content.
    • Top-p: For image generation, a top_p of approximately 0.8 is recommended.
    • Guidance: Increasing cfg_scales slightly during RGB $\rightarrow$ X inference (e.g., surface normal or segmentation) can improve alignment with the input.
    • Super Resolution: Avoid high cfg_scales and low temperatures during super resolution to prevent blurry results.

    SAM Instance Generation

    • For denser estimation of SAM instances, use the generate_sam_dense method (available in notebooks/generation_4M-21.ipynb).
    • Warning: Do not use the output of generate_sam_dense as a conditioning input, as the high token count can cause memory issues.

  12. Perform pre-tokenization of image-like modalities

    main

    To prevent tokenization from becoming a training bottleneck, you can pre-compute tokens for all image-like modalities before training begins. This allows the trainer to load discrete tokens directly.

    How to pre-tokenize: Run the save_vq_tokens.py script with the appropriate arguments for your modality.

    Important Note on Alignment and Cropping: If you are working with non-square images or using --n_crops > 1, pre-tokenization requires cropping. To maintain alignment across modalities, the system automatically creates a crop_settings directory containing crop information for all samples during the first tokenization pass. This directory must be used when tokenizing subsequent modalities of the same dataset to ensure consistent cropping.