Error Aware Distribution Prediction for Lightweight Implicit Neural Representations

Technical Summary: Error Aware Distribution Prediction for Lightweight Implicit Neural Representations

Problem Statement
Implicit Neural Representations (INRs) have emerged as a powerful tool for the compact encoding of scientific data, such as 3D volumes and time-varying scalar fields, offering superior compression ratios compared to traditional lossy methods. However, as lossy approximators, INRs inevitably introduce prediction errors that can compromise the reliability of downstream scientific analysis. While global metrics like Peak Signal-to-Noise Ratio (PSNR) provide an overall error estimate, they fail to capture localized inaccuracies critical for scientific interpretation.

Recent approaches utilize prediction uncertainty as an indicator of relative error scales. However, existing methods face significant limitations:

1. Computational Cost: Classical uncertainty quantification methods, such as deep ensembles and Monte Carlo dropout, introduce substantial computational overhead.
2. Restrictive Assumptions: Lightweight single-pass solutions, such as heteroscedastic and evidential regression, rely on predefined parametric assumptions (e.g., Gaussian or Student's t-distributions). These assumptions limit their ability to model complex scientific data characterized by long-tailed or multi-modal distributions, often resulting in uncertainty fields that fail to align with actual error regions.

Methodology
The authors propose a lightweight training framework that reformulates INR regression as a classification task to enable flexible distribution modeling without restrictive parametric assumptions.

• Discretization Strategy: Continuous target signals are normalized to the range $[0, 1]$ and discretized into $n$ equal-range bins (default $n=256$). The target label for training corresponds to the bin containing the ground-truth value.
• Prediction Mechanism: Instead of predicting a single discrete label, the model outputs a probability distribution over the bins. The final reconstruction value is computed as the expected value of this distribution: $\sum p_i b_i$, where $p_i$ is the probability of bin $i$ and $b_i$ is the bin's midpoint.
• Loss Function: The primary objective is Cross-Entropy (CE) loss, which encourages the model to learn a meaningful probability distribution reflecting the ground truth. To mitigate "salt-and-pepper" noise and quantization errors common in pure classification approaches, a Mean Squared Error (MSE) regularization term is added, anchoring the expected value of the distribution to the ground truth: $L = L_{CE} + \lambda \cdot L_{MSE}$.
• Uncertainty Quantification: Prediction uncertainty is derived from the variance of the predicted probability distribution. This variance serves as an intuitive measure of error awareness, indicating how far a prediction might deviate from its expected value.

Experimental Setup
The study evaluates the proposed method against regression-based baselines (Heteroscedastic Regression and Evidential Regression) across four datasets: two simulation datasets (Rayleigh–Taylor instability, Magnetic Reconnection) and two CT datasets (Bonsai, Aneurism). All models utilize a 5-layer MLP with 256 neurons per layer, ensuring a comparable parameter count across settings.

Key Results

1. Reconstruction Quality: Classification-based approaches (specifically CE and CE+MSE) achieve PSNR scores comparable to or exceeding regression-based methods. On the CT datasets (Bonsai and Aneurism), the classification methods outperform evidential and heteroscedastic regression by significant margins (e.g., up to 5 dB on Bonsai). Qualitative analysis confirms that classification-based INRs better reconstruct high-frequency features, such as thin vessels and small branches, which are often lost or inadequately captured by regression methods constrained by Gaussian assumptions.
2. Error Awareness (AUSE): The Area Under the Sparsification Error (AUSE) metric is used to evaluate how well uncertainty aligns with true prediction errors.
   • Cross-Entropy (CE) and Evidential Regression generally achieve the lowest AUSE scores.
   • Pure MSE regression yields the highest (worst) AUSE scores, as it does not explicitly optimize for meaningful probability distributions.
   • The addition of MSE regularization to CE (CE+MSE) slightly increases AUSE compared to pure CE but remains competitive, offering a balance between reconstruction accuracy and uncertainty quality.
3. Distribution Flexibility: Visualizations of predicted distributions at sharp gradient regions demonstrate that the classification setting can capture multi-modal behaviors, whereas heteroscedastic regression is forced into a unimodal Gaussian shape, failing to represent complex error structures.

Significance and Claims
The paper claims that reformulating INR training as a classification task offers a practical and effective alternative for scientific data modeling. The primary significance lies in achieving a complementary performance profile:

• High Reconstruction Quality: The method preserves fine-grained, high-frequency features better than parametric regression approaches.
• Meaningful Uncertainty: By avoiding restrictive distributional assumptions, the classification approach generates uncertainty fields that align more closely with actual error fields, particularly in capturing localized inaccuracies.

The authors conclude that while a trade-off exists between reconstruction accuracy and uncertainty estimation across different methods, the classification-based framework provides a lightweight, single-pass solution that effectively balances both requirements, making it suitable for scientific workflows where both data fidelity and error awareness are critical. Future work is proposed to investigate the theoretical reasons why this framework captures error awareness in deterministic data and to explore its application in settings with inherent uncertainty sources, such as ensemble simulations.