Local-Order Auxiliary Losses Can Improve Autoencoder Reconstruction

Technical Summary: Local-Order Auxiliary Losses Can Improve Autoencoder Reconstruction

Problem Statement
In autoencoders designed to reconstruct numerical arrays (tensors), the standard training objective is Mean Squared Error (MSE). While MSE ensures pointwise accuracy, compressed reconstructions often depend on preserving local spatial order (e.g., the relative arrangement of ridges, troughs, and fronts). Conventional wisdom suggests a trade-off: optimizing for structural metrics (like structural similarity or gradient fidelity) often degrades pointwise MSE. This paper investigates whether this trade-off is necessary. Specifically, it asks if structural auxiliary losses can improve validation MSE in finite-capacity autoencoders, rather than merely improving perceptual metrics at the expense of pointwise accuracy.

Methodology
The authors introduce Finite-Difference Sign Error (FDSE), a local-order auxiliary objective designed to penalize disagreements between the signs of neighboring finite differences in the target and the reconstruction.

• Definition: For a rank-$d$ tensor $X$, FDSE computes forward differences along each axis. It replaces exact sign comparison with a differentiable approximation using a smooth, odd, monotone, saturating function $\phi_s$ (e.g., scaled arctan or tanh). The loss is defined as:
  $$\text{FDSE}_{\phi,s}(X, \hat{X}) = \sum_{i=1}^d \| \phi_s(\delta_i(X)) - \phi_s(\delta_i(\hat{X})) \|_1$$
  where $\delta_i$ denotes the forward difference along axis $i$.
• Training Objective: The model is trained using a convex mixture of MSE and FDSE:
  $$\mathcal{L}_{\text{train}}(\theta) = \lambda \text{MSE}(X, A_\theta(X)) + \mu \text{FDSE}_{\phi,s}(X, A_\theta(X))$$
  where $\lambda > 0$ and $\mu \ge 0$.
• Rationale: FDSE is invariant to positive affine transformations (offset and scale), meaning it cannot control amplitude or offset on its own; thus, it must be paired with MSE. The authors argue that in capacity-limited settings, a local-order error (e.g., flipping the sign of a finite difference) often implies a displaced or flattened feature, which inherently increases pointwise error. By correcting local ordering, the model may inadvertently reduce pointwise MSE.
• Experimental Setup: The study evaluates four reconstruction tasks: shallow-water states, ERA5 potential-vorticity blocks, PDEBench radial-dam-break scalar fields, and CIFAR-10 images. Four different autoencoder architectures are used (CRA5, NDAE, MBT-style ND codec, and a factorized-prior image autoencoder). The authors compare MSE-only training against mixtures of MSE with FDSE and seven other standard auxiliary objectives (L1, gradient magnitude, SSIM, Laplacian residual, Total Variation, and spectral magnitude).

Key Contributions

1. Coupled Improvement: The paper demonstrates that in compressed-latent reconstruction, structural auxiliary losses can simultaneously improve non-pointwise structural metrics and validation MSE, challenging the assumption that these objectives must trade off against each other.
2. FDSE Objective: The introduction of FDSE as a lightweight, architecture-agnostic, differentiable local-order supervision that penalizes sign mismatches in adjacent finite differences without requiring exact derivative magnitude matching.
3. Empirical Evidence: Across four distinct tasks and autoencoder pairs, moderate mixtures of MSE and FDSE consistently achieve lower validation MSE than pure MSE training.
4. Diagnostic Analysis: The authors employ "projection-matched" controls to rule out the hypothesis that the improvement is merely an implicit step-size effect (i.e., that the auxiliary loss simply rescales the MSE gradient). They show that the FDSE gain includes a directionally distinct update component that pure MSE training, even with adjusted learning rates, cannot replicate.

Results

• MSE Reduction: In coefficient sweeps across the four tasks, moderate MSE:FDSE mixtures reduced validation MSE by factors ranging from 2.3$\times$ to 7.0$\times$ compared to pure MSE training.
  • Example: On CIFAR-10, validation MSE dropped from $\approx 3.0 \times 10^{-3}$ (pure MSE) to $\approx 4.3 \times 10^{-4}$ (0.8:0.2 mixture).
  • Example: On ERA5 potential-vorticity blocks, MSE was reduced from $4.60 \times 10^{-4}$ to $1.02 \times 10^{-4}$ (4.5$\times$ reduction) with a 0.6:0.4 mixture.
• Structural Metrics: Improvements in MSE were generally accompanied by improvements in topology-sensitive diagnostics (PERSD, MERGD, CORRD) and SSIM.
• Comparison with Other Auxiliaries: While other auxiliaries (like L1 or SSIM) sometimes improved specific metrics, FDSE was frequently among the strongest structural objectives tested. Notably, pure FDSE performed poorly, confirming the necessity of the MSE term for anchoring amplitude and offset.
• Applicability: The gains were most pronounced in tasks with coherent spatial fields (ERA5, shallow water, PDEBench) where local order carries significant information. The method is less effective for texture-dominated signals or axes stacking unrelated physical quantities.

Significance and Claims
The paper claims that appropriately weighted local-structure supervision can guide optimization toward solutions with better pointwise accuracy in compressed-latent reconstruction, rather than merely improving perceptual metrics at MSE's expense. The authors emphasize that this is not a universal rule; the effect is contingent on the target field having coherent local variation at the resolved scale and the training objective being fixed-target reconstruction under a value-fidelity term.

The work suggests that the conventional trade-off between structural diagnostics and MSE is often inverted in these specific settings. By treating the mixture weight as a validation-selected hyperparameter rather than proposing a complex dynamic balancing rule, the authors show that simple, moderate mixtures of MSE and FDSE can yield superior reconstructions for scientific and analytical tensors, potentially improving downstream visualization and analysis. However, they caution that using FDSE outside its applicability conditions (e.g., on texture-dominated data) could preserve incorrect local structures and distort application-relevant features.