Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance (2306.02866v3)

Abstract: We present a novel framework to overcome the limitations of equivariant architectures in learning functions with group symmetries. In contrary to equivariant architectures, we use an arbitrary base model such as an MLP or a transformer and symmetrize it to be equivariant to the given group by employing a small equivariant network that parameterizes the probabilistic distribution underlying the symmetrization. The distribution is end-to-end trained with the base model which can maximize performance while reducing sample complexity of symmetrization. We show that this approach ensures not only equivariance to given group but also universal approximation capability in expectation. We implement our method on various base models, including patch-based transformers that can be initialized from pretrained vision transformers, and test them for a wide range of symmetry groups including permutation and Euclidean groups and their combinations. Empirical tests show competitive results against tailored equivariant architectures, suggesting the potential for learning equivariant functions for diverse groups using a non-equivariant universal base architecture. We further show evidence of enhanced learning in symmetric modalities, like graphs, when pretrained from non-symmetric modalities, like vision. Code is available at https://github.com/jw9730/lps.

• The paper introduces a probabilistic symmetrization method that transforms any neural architecture into an equivariant model.
• It employs an input-dependent parameterization of group averaging to achieve universal approximation with reduced sample complexity.
• Empirical tests in graph classification and particle dynamics demonstrate competitive performance against tailored equivariant networks.

Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance

This paper introduces a framework for achieving group equivariance in neural networks using a concept termed "probabilistic symmetrization." The paper addresses the limitations inherent in existing equivariant architectures—such as restricted expressive power and architectural rigidity—by proposing a method that utilizes an arbitrary base model, including MLPs or transformers, and symmetrizes it to achieve equivariance concerning a specified group.

The core idea differs from traditional approaches that construct models explicitly constrained by group equivariance. Instead, the proposed framework employs a smaller auxiliary network to parameterize a probabilistic distribution governing the symmetrization process. This distribution is trained end-to-end alongside the base model, optimizing performance while reducing the sample complexity inherent in symmetrization. An essential aspect of the framework is ensuring both group equivariance and universal approximation capacity in expectation.

Methodology

The paper's proposed method for achieving equivariant architecture-agnostic networks involves replacing the uniform distribution typically used in group averaging with a parameterized, input-dependent probabilistic distribution $p_\omega(g|\mathbf{x})$. This distribution is conditioned on the input data and symmetrizes the base function by following:

$$\phi_{\theta,\omega}(\mathbf{x}) = \mathbb{E}_{p_\omega(g|\mathbf{x})}\left[\rho_2(g)f_\theta(\rho_1(g)^{-1}\mathbf{x})\right],$$

where the distribution needs to meet conditions of probabilistic equivariance. By incorporating a noise-outsourced map with a small equivariant network, this distribution coordinates with the base function and maintains equivariance and expressive power vanishing the need for group-specific model design.

The authors simultaneously implement this approach across various core models, including MLP and transformers initialized from pre-trained vision transformers. These models were evaluated on numerous symmetry groups like permutation and Euclidean groups, where empirical results confirmed competitive performance against tailored equivariant architectures.

Experimental Findings

The experimental setup spanned multiple domains, reflecting the generality of their approach.

1. Graph Isomorphism Learning: The framework outperformed existing MLP and equivariant network approaches by achieving high graph classification accuracy on tasks like GRAPH8c and EXP-classify, where the sample complexity is notably challenging.
2. Particle Dynamics Learning: On the $n$-body problem dataset, the framework demonstrated superior capability in predicting particle dynamics, exceeding tailored equivariant networks by minimizing the Mean Square Error (MSE) on position prediction.
3. Pattern Recognition in Graphs: The paper further validates the framework on a node classification task (PATTERN dataset), illustrating that there is an additional boost in performance when the base model benefits from transfer learning via pre-trained weights from other modalities, such as vision transformers.
4. Real-World Graph Tasks: Deploying a vision transformer symmetrized with their method onto real-world graph data, the research showcased improvements over baselines, particularly in domains requiring nuanced understanding like molecular and peptide properties.

Broader Implications and Future Work

This work opens promising avenues for neural networks by facilitating general-purpose, equivariant learning across various symmetry groups without heavily relying on specialized architectures. These findings suggest practical applications across fields such as computational chemistry, physics, and biology, where the symmetry of the problem structure is prevalent but computational efficiency and transferability are crucial.

The expectation is that future research will pivot on enhancing parameter sharing and knowledge transfer across different domains using the proposed symmetrization method. Additionally, advancing computational efficiency during inference remains an issue, notably with the necessity of estimating outcomes with multiple samples, pointing towards potential gains in algorithmic optimization or parallel processing techniques.

In sum, this paper presents a strong case for using probabilistic symmetrization to unlock architecture-agnostic equivariance, heralding more flexible and transferable neural network models. The authors hint at continued investigation into parameter efficiency, potentially paving the way for even broader application and adoption in complex, symmetry-rich domains.