Supplementary A Model Implementation Details

During training, we also apply a random translation augmentation to all the particle positions to artificially generate more training data. Next, we describe the KL variant of this model. Note that for the cosmology example in section 4.3, we use the L