Atomium90/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

Category	Variants
Positional encoding	RoPE, ALiBi vs learned absolute
Feed-forward block	SwiGLU vs GELU
Weight initialization	He, Xavier, DeepNorm vs normal
Model depth	depth 4 / 8 / 12 / 16 vs baseline 6
Residual scaling	fixed α, depth-scaled α, ReZero vs none

Key results (val loss, lower is better)

Model	Params	Loss	Δ baseline	FID
baseline	109.86M	3.499	—	47.2
swiglu_v1	122.44M	3.446	-0.053	59.7
rope_v1	109.86M	3.492	-0.007	41.3
swiglu_depth16	216.85M	3.313	-0.186	38.9
swiglu_depth16_rope_deepnorm_residual_depth	216.85M	3.286	-0.213	29.2

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