diffusion/README.md

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.