DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows
This work addresses the challenge of applying DEQs to permutation-invariant data like point clouds, offering a novel extension with potential benefits in efficiency for domain-specific applications.
The paper tackled the problem of extending Deep Equilibrium Models (DEQs) to handle discrete measure inputs like sets or point clouds, by introducing Distributional Deep Equilibrium Models (DDEQs) with a theoretical framework based on Wasserstein gradient flows, and showed that DDEQs can compete with state-of-the-art models in tasks such as point cloud classification and completion while being more parameter-efficient.
Deep Equilibrium Models (DEQs) are a class of implicit neural networks that solve for a fixed point of a neural network in their forward pass. Traditionally, DEQs take sequences as inputs, but have since been applied to a variety of data. In this work, we present Distributional Deep Equilibrium Models (DDEQs), extending DEQs to discrete measure inputs, such as sets or point clouds. We provide a theoretically grounded framework for DDEQs. Leveraging Wasserstein gradient flows, we show how the forward pass of the DEQ can be adapted to find fixed points of discrete measures under permutation-invariance, and derive adequate network architectures for DDEQs. In experiments, we show that they can compete with state-of-the-art models in tasks such as point cloud classification and point cloud completion, while being significantly more parameter-efficient.