| # 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. |
|
|
