Official PyTorch implementation of Spectral Convolutional Conditional Neural Processes (SConvCNP), accepted at NeurIPS 2025.

📄 Paper:
Spectral Convolutional Conditional Neural Processes

This repository contains the implementation code for training and evaluating spectral convolutional conditional neural processes and baseline models on various regression tasks. This codebase is kept simple and straightforward to facilitate its use for research purposes.

• Python 3.11+ (tested with 3.11.5 and 3.13.1)
• PyTorch 2.0+
• CUDA-enabled GPU (strongly recommended for training)

Note: The environment specifications in ENVIRONMENT.md document the setup used for code development and debugging on macOS. For reproducing paper results, use CUDA-enabled GPU servers as described in the Training Environment section of ENVIRONMENT.md.

SConvCNP/
├── nps/                          # Core model implementations
│   ├── models/                   # Model architectures (CNP, ConvCNP, TNP, etc.)
│   ├── core/                     # Core components
│   │   ├── cnns/                 # CNN architectures (UNet, FNO, UNO)
│   │   ├── encoders/             # Encoder modules
│   │   ├── decoders/             # Decoder modules
│   │   ├── convolutions/         # Convolution operations
│   │   └── ...
│   └── likelihoods/              # Likelihood functions
├── utils/                        # Utilities
│   ├── data/                     # Data generators and datasets
│   │   ├── synthetic/            # Synthetic data (GP, sawtooth, etc.)
│   │   ├── image/                # Image datasets (DTD)
│   │   ├── predprey/             # Predator-Prey dynamics
│   │   └── traffic/              # Traffic forecasting
│   ├── experiment/               # Experiment utilities
│   └── plot_fn/                  # Plotting functions
├── scripts/                      # Configuration files
│   ├── synthetic/configs/        # Synthetic experiment configs
│   ├── dtd/configs/              # DTD dataset configs
│   ├── predprey/configs/         # Predator-Prey configs
│   └── traffic/configs/          # Traffic forecasting configs
├── train.py                      # Training script
├── eval.py                       # Evaluation script
└── README.md                     # This file

The following models are implemented:

• CNP - Conditional Neural Process
• ACNP - Attentive Conditional Neural Process
• ConvCNP - Convolutional Conditional Neural Process (with UNet/ResNet)
• SConvCNP - Spectral Convolutional CNP (with FNO/UNO/UNO2)
• TNP - Transformer Neural Process
• TE-TNP - Translation Equivariant TNP

• Gaussian Processes (RBF, Matern-1/2, Matern-3/2, Matern-5/2, Periodic kernels)
• Sawtooth functions
• Square wave functions

• DTD (Describable Textures Dataset) - Image completion task
• Predator-Prey - Lotka-Volterra dynamics
• Traffic - Traffic forecasting

The training script uses a hierarchical YAML configuration system. You must specify:

1. A base configuration file for the dataset
2. Additional configuration files for data generators (synthetic only) and model architecture

Sawtooth function with SConvCNP-UNO:

python train.py --config \
    scripts/synthetic/configs/base.yaml \
    scripts/synthetic/configs/data/output_generators/sawtooth.yaml \
    scripts/synthetic/configs/data/input_generators/uniform.yaml \
    scripts/synthetic/configs/models/sconvcnp-uno.yaml

Gaussian Process (Matern-1/2 kernel) with ConvCNP-UNet:

python train.py --config \
    scripts/synthetic/configs/base.yaml \
    scripts/synthetic/configs/data/output_generators/gp-matern12.yaml \
    scripts/synthetic/configs/data/input_generators/uniform.yaml \
    scripts/synthetic/configs/models/convcnp-unet.yaml

Other output generators: gp-rbf.yaml, gp-matern32.yaml, gp-matern52.yaml, gp-periodic.yaml, squarewave.yaml

python train.py --config \
    scripts/dtd/configs/base.yaml \
    scripts/dtd/configs/models/sconvcnp-uno.yaml

Available models: cnp.yaml, acnp.yaml, convcnp-resnet.yaml, sconvcnp-fno.yaml, sconvcnp-uno.yaml, sconvcnp-uno2.yaml, eqtnp.yaml, te-eqtnp.yaml

python train.py --config \
    scripts/predprey/configs/base.yaml \
    scripts/predprey/configs/models/sconvcnp-uno.yaml

Available models: cnp.yaml, acnp.yaml, convcnp-unet.yaml, sconvcnp-uno.yaml, sconvcnp-uno2.yaml, eqtnp.yaml, te-eqtnp.yaml

python train.py --config \
    scripts/traffic/configs/base.yaml \
    scripts/traffic/configs/models/sconvcnp-uno.yaml

Key training parameters can be modified in the base configuration files:

• misc.epochs: Number of training epochs (default: 500)
• misc.seed: Random seed for reproducibility (default: 0)
• optimizer.lr: Learning rate (varies by dataset)
• misc.gradient_clip_val: Gradient clipping value (default: 0.5)
• misc.wandb_logging_enabled: Enable/disable Weights & Biases logging (default: true)

Checkpointing is configured in the misc.checkpointing section:

• Enable/Disable: Set misc.checkpointing.enabled to true/false
• Local Checkpoints: Saved to artifacts/sconvcnp/{experiment}/... when misc.checkpointing.local is true
• Resume Training: Set misc.checkpointing.resume.enabled: true and provide checkpoint_path

Checkpoints are saved in:

artifacts/sconvcnp/{experiment_name}/x={x_dim}d_y={y_dim}d/model={model_name}/seed={seed}/checkpoints/

The evaluation script loads trained checkpoints and evaluates them on the test set.

python eval.py --wandb-run-id <wandb_run_id> --config <config_files...>

Important: You must provide the same configuration files used during training.

Note: The --wandb-run-id parameter is optional. If not provided, evaluation will use the checkpoint specified in the configuration files or the most recent checkpoint available.

python eval.py --wandb-run-id <entity/project/run_id> --config \
    scripts/synthetic/configs/base.yaml \
    scripts/synthetic/configs/data/output_generators/sawtooth.yaml \
    scripts/synthetic/configs/data/input_generators/uniform.yaml \
    scripts/synthetic/configs/models/sconvcnp-uno.yaml

python eval.py --wandb-run-id <entity/project/run_id> --config \
    scripts/dtd/configs/base.yaml \
    scripts/dtd/configs/models/sconvcnp-uno.yaml

python eval.py --wandb-run-id <entity/project/run_id> --config \
    scripts/predprey/configs/base.yaml \
    scripts/predprey/configs/models/sconvcnp-uno.yaml

The evaluation script will:

1. Load the trained checkpoint
2. Evaluate on the test set
3. Log metrics (log-likelihood, RMSE) to console and Weights & Biases
4. Generate visualization plots (if enabled)

The codebase uses a hierarchical YAML configuration system powered by HiYaPyCo and Hydra. Configuration files are composed of:

1. Base config - Dataset-specific settings (data generators, plotting, artifacts paths)
2. Data generator configs - Input/output distributions (synthetic only)
3. Model configs - Model architecture and hyperparameters

Configuration files use variable interpolation (e.g., ${params.x_dim}) and support evaluation of expressions (e.g., ${eval:'2 * ${.resolution}'}).

By default, experiments are logged to Weights & Biases. Configure in base config:

misc:
  wandb_logging_enabled: true
  wandb_user: <your-username>
  project: SConvCNP-{experiment_name}

Disable wandb logging:

misc:
  wandb_logging_enabled: false

Training artifacts are saved to:

artifacts/sconvcnp/{experiment}/
├── checkpoints/       # Model checkpoints
├── metrics/           # Metrics logs
├── plots/             # Visualization plots
└── train/wandb/       # Wandb logs

If you use this code in your research, please cite:

@article{mohseni2026spectral,
  title={Spectral convolutional conditional neural processes},
  author={Mohseni, Peiman and Duffield, Nick},
  journal={Advances in Neural Information Processing Systems},
  volume={38},
  pages={40416--40446},
  year={2026}
}

This codebase adapts ideas and code from several open-source projects, which we gratefully acknowledge:

• TETNP — Translation-Equivariant Transformer Neural Processes (Cambridge MLG): the translation-equivariant transformer / attention layers and the TE-TNP model family.
• neuralprocesses (Wessel Bruinsma): convolutional NP building blocks (UNet / ConvBlock coders) and the predator–prey (Lotka–Volterra → Hudson Bay hare–lynx) data pipeline.

Each project remains under its own license; please consult the linked repositories.

This project is released under the MIT License.