# Conditional diffusion models, 200 epochs with classifier-free guidance

The six conditional diffusion models used for the figures of the NeurIPS 2026 paper "A multi-scale information geometry reveals the structure of mutual information in neural populations". There is one model per encoder (Inception, ResNet-50, ON-OFF) and area (V1, V4).

## Contents

    inception/V1/final_model.pt     inception/V4/final_model.pt
    resnet50/V1/final_model.pt      resnet50/V4/final_model.pt
    onoff/V1/final_model.pt         onoff/V4/final_model.pt
    */*/train_log.json              loss per epoch
    diffusion.yaml                  training configuration, shared by all six models
    SHA256SUMS                      checksums of the checkpoints

## Training recipe

The recipe is the same for all six models:

- Network: a single-channel 64x64 DDPM U-Net with block widths (64, 128, 256, 512) and two layers per block, conditioned on the 50-dimensional response vector through an MLP token (dimension 256) and cross-attention.
- Training data: 100,000 TinyImageNet training images, converted to grayscale and resized to 64x64, each paired with the responses predicted by the encoder for that area, with observation noise of sigma = 0.22 added at every step.
- Classifier-free guidance: the responses are dropped (null token) with probability 0.1, so the same network also gives the unconditional prediction.
- Noise schedule: 1000 steps, linear schedule with beta from 1e-4 to 0.02, epsilon prediction.
- Optimiser: AdamW with learning rate 1e-4 and weight decay 1e-4, cosine schedule with 500 warm-up steps, batch size 128, bfloat16 precision, gradient clipping at 1.0, seed 42.
- Length: 200 epochs.

## Format

Each `final_model.pt` is a dictionary with the keys:

- `model_state_dict`: the network weights;
- `epoch`: 200;
- `loss`: the final training loss;
- `config`: the configuration file and the response file used for training;
- `encoder_name` and `area`: which model this is;
- `provenance`: the git commit of the training code and the command line, with paths made relative for this release.