Symmetry-Aware Generative Modeling through Learned Canonicalization

2025-01-14

AI Summary

Modeling density distributions of symmetric structures—such as molecules—remains challenging due to inherent permutation and rotational symmetries, which conventional equivariant generative models struggle to handle efficiently. Method: This paper introduces a novel non-equivariant diffusion modeling paradigm for symmetric data: instead of enforcing group equivariance in the generative model, it first learns a group-equivariant normalization network to map each input to a unique canonical pose, then trains a standard (non-equivariant) diffusion model exclusively in this canonical space. Contribution/Results: This decouples symmetry handling into two distinct stages—normalization followed by non-equivariant modeling—thereby avoiding the representational limitations, parameter redundancy, and optimization difficulties of fully equivariant architectures. Crucially, it ensures that each group orbit is represented by exactly one canonical sample. On molecular conformation generation, the method achieves superior sampling quality and diversity, while accelerating inference by 2.3×, demonstrating significant gains in efficiency, accuracy, and practicality.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Non-convex Optimization

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Generative modeling of symmetric densities has a range of applications in AI for science, from drug discovery to physics simulations. The existing generative modeling paradigm for invariant densities combines an invariant prior with an equivariant generative process. However, we observe that this technique is not necessary and has several drawbacks resulting from the limitations of equivariant networks. Instead, we propose to model a learned slice of the density so that only one representative element per orbit is learned. To accomplish this, we learn a group-equivariant canonicalization network that maps training samples to a canonical pose and train a non-equivariant generative model over these canonicalized samples. We implement this idea in the context of diffusion models. Our preliminary experimental results on molecular modeling are promising, demonstrating improved sample quality and faster inference time.

Problem

• Symmetry-aware Models
• Molecular Structure
• Drug Discovery

Innovation

• Symmetry Learning
• Diffusion Model
• Molecular Modeling