Explain the implementation-dependent performance discrepancy between Linear and MetaLayer Deep Set architectures

Determine the mechanism by which PyTorch Linear-based and PyTorch Geometric MetaLayer-based implementations of permutation-invariant Deep Set models produce different optimization behavior and predictive performance when inferring the cosmological matter-density parameter Ω_m from galaxy line-of-sight peculiar velocities.

Background

The paper compares two functionally similar set-based implementations for inferring Ω_m from galaxy peculiar velocities: an MLP-based Deep Set constructed with standard PyTorch Linear layers and a MetaLayer-based implementation from the PyTorch Geometric library. Although both process galaxies independently and aggregate their representations permutation-invariantly, the MetaLayer implementation performs substantially worse in the reported experiments.

The authors localize the discrepancy to the layer architecture rather than to the input features or training simulation suite, but do not identify why the implementations exhibit different optimization behavior and predictive accuracy. They characterize the comparison as provisional and indicate that a more detailed diagnostic study is needed.