Code for paper "Multi-Fidelity Residual Neural Processes for Scalable Surrogate Modeling", presented at ICML2024.

Multi-fidelity surrogate modeling aims to learn an accurate surrogate at the highest fidelity level by combining data from multiple sources. Traditional methods relying on Gaussian processes can hardly scale to high-dimensional data. Deep learning approaches utilize neural network based encoders and decoders to improve scalability. These approaches share encoded representations across fidelities without including corresponding decoder parameters. This hinders inference performance, especially in out-of-distribution scenarios when the highest fidelity data has limited domain coverage. To address these limitations, we propose Multi-fidelity Residual Neural Processes (MFRNP), a novel multi-fidelity surrogate modeling framework. MFRNP explicitly models the residual between the aggregated output from lower fidelities and ground truth at the highest fidelity. The aggregation introduces decoders into the information sharing step and optimizes lower fidelity decoders to accurately capture both in-fidelity and crossfidelity information. We show that MFRNP significantly outperforms state-of-the-art in learning partial differential equations and a real-world climate modeling task.

MFRNP Model Structure

Create conda environment and install the packages:

conda env create -f environment.yml

Unzip the file (data.zip) at the root directory.

unzip data.zip

Run Full and OOD tasks (Fluid, Heat2,3,5 and Poisson2,3,5):

./run_task.sh

Run ablation study with MFRNP-H:

./run_task_ablation.sh

Results are saved at "result" directory.

python train.py --data_path <path_to_dataset> --save_dir <name_of_directory_to_be_saved_in_result_folder> --config <path_to_config_file> --levels <#_of_total_fidelity_levels> --device <cuda_or_cpu>

python train.py --data_path "data/full_dataset/poisson" --save_dir poisson --config pde_config.yaml --levels 2 --device cuda

@inproceedings{niu2024multi,
  author       = {Niu, Ruijia and Wu, Dongxia and Kim, Kai and Ma, Yi-An and Watson-Parris, Duncan and Yu, Rose},
  title        = {Multi-Fidelity Residual Neural Processes for Scalable Surrogate Modeling},
  booktitle    = {International Conference on Machine Learning, {ICML} 2024},
  series       = {Proceedings of Machine Learning Research},
  year         = {2024}
}