Eq-WaLa: Equivariant Augmentation and Regularization for Wavelet Latent Flow Matching

Despite their strong performance in compact 3D shape modeling, latent implicit generative models typically neglect geometric symmetries of 3D shapes, particularly rotation invariance, causing representations to depend on arbitrary coordinate frames rather than intrinsic shape structure. To tackle this, we propose an equivariant generative framework for 3D shapes that explicitly accounts for rotational variability during training. Our method builds on a compact wavelet-tree representation of a shape's signed distance field (SDF), enabling multi-scale, channel-efficient encoding and efficient convolutional processing. To address the common challenge of unaligned orientations in real-world datasets, we introduce two key components: (i) a novel wavelet-domain rotation augmentation scheme that enables exact on-the-fly construction of rotated inputs via the inherent equivariance of wavelet transforms, and (ii) a latent-space regularization strategy that enforces consistency under discrete latent rotations. Combined with a flow matching model trained on these augmented and regularized latents, our framework can produce high-quality 3D shapes. We validate the effectiveness of these components, showing improved reconstruction accuracy and faster generative convergence on the Thingi10K dataset, which lacks consistent canonical orientations. Furthermore, we also demonstrate the scalability of our approach to large-scale settings, including both unconditional and text-conditioned shape generation tasks on the diverse and complex Objaverse dataset.