  Modeling sets is an important problem in machine learning since this type of
data can be found in many domains. A promising approach defines a family of
permutation invariant densities with continuous normalizing flows. This allows
us to maximize the likelihood directly and sample new realizations with ease.
In this work, we demonstrate how calculating the trace, a crucial step in this
method, raises issues that occur both during training and inference, limiting
its practicality. We propose an alternative way of defining permutation
equivariant transformations that give closed form trace. This leads not only to
improvements while training, but also to better final performance. We
demonstrate the benefits of our approach on point processes and general set
modeling.
