FA-INR: Adaptive Implicit Neural Representations for Interpretable Exploration of Simulation Ensembles

Abstract

Surrogate models are essential for efficient exploration of large-scale ensemble simulations. Implicit neural representations (INRs) provide a compact and continuous framework for modeling spatially structured data, but they often struggle with learning complex localized structures within the scientific fields. Recent INR-based surrogates address this by augmenting INRs with explicit feature structures, but at the cost of flexibility and substantial memory overhead. In this paper, we present Feature-Adaptive INR (FA-INR), an adaptive INR-based surrogate model for high-fidelity and interpretable exploration of ensemble simulations. Instead of relying on structured feature representations, FA-INR leverages cross-attention over a learnable key-value memory bank to allocate model capacity adaptively based on the data characteristics. To further improve scalability, we introduce a coordinate-guided mixture of experts (MoE) framework that enhances both efficiency and specialization of feature representations. More importantly, the learned experts produce an interpretable partition over the simulation domain, enabling scientists to identify complex structures and perform localized parameter-space exploration. Beyond quantitative and qualitative evaluations, we also demonstrate that our learned expert specialization can reveal meaningful scientific insights and support localized sensitivity analysis.

Figures (10)

• Figure 1: Architecture of the proposed FA-INR. An input pair $(x,p)$ is first routed by a gating network in (a) to the Top-2 relevant spatially-specialized encoder experts (see (b)). Then, the aggregated feature vector from all selected experts is further processed by the MLP decoder in (b).
• Figure 2: A single memory bank already outperforms the feature grid under the same feature storage budget. By introducing multiple encoder experts (MoE), our method further scales performance more effectively than simply expanding the memory bank, while requiring fewer feature vectors compared to the grid-based approach.
• Figure 3: Key usage comparison between a single memory bank and the MoE design. In this example, the MoE setup with four experts distributes queries more evenly across keys within each expert, resulting in more effective key utilization.
• Figure 4: The comparison of visual fidelity on a test member of the MPAS-Ocean dataset.
• Figure 5: The comparison of visual fidelity on a test member of the Nyx dataset. The PSNR reported below each image denotes the reconstruction accuracy of the underlying scalar field.