nano4M-architecture-ablations

nano4M — Architecture Improvements

COM-304 Foundation Models Project — EPFL Spring 2026
Edouard Fousson · Thomas Picart · Jonathan Balli

Code & training scripts: github.com/Atomium90/nano4M-architectural-improvements

Systematic ablation of five architectural modifications to the nano4M baseline transformer, trained on multimodal CLEVR (RGB, Depth, Surface Normals, Scene Descriptions).

Modifications studied

Key results (val loss, lower is better)

See results_summary.csv for the full table including compute-fair evaluations.

Compute-fair evaluation

Models with more parameters use more FLOPs per step. For fair comparison, we additionally evaluate each variant at a checkpoint matching the baseline's total training compute:

fair_tokens = 5000M × (N_baseline / N_variant)

Each experiment folder contains both checkpoint-final.safetensors (full 5000M token training) and the compute-fair intermediate checkpoint where applicable.

Usage

import yaml, omegaconf, torch
from hydra.utils import instantiate
from safetensors.torch import load_file

def load_model(checkpoint_path, config_path, device="cuda"):
    with open(config_path) as f:
        raw = yaml.safe_load(f)
    omegaconf.OmegaConf.register_new_resolver("eval", eval, replace=True)
    cfg = omegaconf.OmegaConf.to_container(omegaconf.OmegaConf.create(raw), resolve=True)
    model = instantiate(cfg["model_config"]).to(device)
    model.load_state_dict(load_file(checkpoint_path, device=device), strict=True)
    return model.eval()

# Example
model = load_model(
    "swiglu_depth16/checkpoint-final.safetensors",
    "configs/variants/swiglu_depth16.yaml",
)

Configs for all variants are in the configs/ folder of this repo.

Training setup

• Dataset: pre-tokenized multimodal CLEVR (Cosmos DI16x16 tokenizer)
• Total tokens: 5 000M per run
• Batch size: 512 (constant across all runs via batch_size × num_gpus = 512)
• Optimizer: AdamW, lr=6e-4, cosine schedule
• Hardware: 2–4× NVIDIA L40S or 2× H100

Repository structure

├── <exp_name>/
│   ├── checkpoint-final.safetensors        ← full training
│   └── checkpoint-<step>.safetensors       ← compute-fair (if applicable)
├── configs/
│   ├── multiclevr_d6-6w512.yaml            ← baseline config
│   └── variants/                           ← one YAML per experiment
└── results_summary.csv                     ← all metrics