File size: 2,105 Bytes
556cdf0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | # 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.
|