Interpreting Equivariant Representations

TL;DR

This work addresses how latent representations learned under geometric inductive biases should be interpreted when the data are acted on by transformation groups. It introduces invariant projections of latent spaces to obtain unambiguous, Euclidean representations and demonstrates that these invariant spaces can reveal structure and preserve information better than the raw equivariant representations. The authors formalize the relationship between equivariant latent codes and quotient spaces, propose conditions for isometric cross sections and, when unavailable, random invariant projections, and validate the approach on a permutation-equivariant graph VAE for QM9 and a rotation-equivariant MNIST classifier. The results show that invariant representations yield clearer structure, more stable interpolations, and improved downstream performance, providing practical tools for developers and researchers to analyze and use equivariant representations beyond training.

Abstract

Latent representations are used extensively for downstream tasks, such as visualization, interpolation or feature extraction of deep learning models. Invariant and equivariant neural networks are powerful and well-established models for enforcing inductive biases. In this paper, we demonstrate that the inductive bias imposed on the by an equivariant model must also be taken into account when using latent representations. We show how not accounting for the inductive biases leads to decreased performance on downstream tasks, and vice versa, how accounting for inductive biases can be done effectively by using an invariant projection of the latent representations. We propose principles for how to choose such a projection, and show the impact of using these principles in two common examples: First, we study a permutation equivariant variational auto-encoder trained for molecule graph generation; here we show that invariant projections can be designed that incur no loss of information in the resulting invariant representation. Next, we study a rotation-equivariant representation used for image classification. Here, we illustrate how random invariant projections can be used to obtain an invariant representation with a high degree of retained information. In both cases, the analysis of invariant latent representations proves superior to their equivariant counterparts. Finally, we illustrate that the phenomena documented here for equivariant neural networks have counterparts in standard neural networks where invariance is encouraged via augmentation. Thus, while these ambiguities may be known by experienced developers of equivariant models, we make both the knowledge as well as effective tools to handle the ambiguities available to the broader community.

Figures (10)

• Figure 1: Visualizing rotated MNIST images via their equivariant representation (left) hides structure that is apparent in an invariant representation of the same latent codes (right).
• Figure 2: An illustration of the effect of applying sorting to $\mathbb{R}^2$. In this case, we see, that $\mathbb{R}^2$ is mapped to $\mathcal{Z}_s = \{(x,y) \in \mathbb{R}^2 | x \leq y\}$ (i.e. the blue area). We see, that in the $S_2$ equivariant representation depicted on the left-hand side the similar objects (e.g. the pentagons) are not guaranteed to be close. After having applied an invariant map (sorting), we see, that this shortcoming is taken care of.
• Figure 3: An illustration of possible properties of the induced map $s': \mathcal{Z}/G \to \mathcal{Z}_s$ to prioritize when choosing an invariant projection $s: \mathcal{Z} \to \mathcal{Z}_s$. Each category of mappings preserves a different amount of structure of the latent space.
• Figure 4: The first two principal components of the QM9 training dataset for the equivariant (top) and invariant representation (bottom) of the latent space. Each column illustrates a specific molecular property.
• Figure 5: Molecules generated from interpolating between two molecules using the equivariant (top) and invariant (bottom) representations. Note that while the molecules decoded from $z_2$ and $s(z_2)$ differ in their embedding, they are equal up to permutation. Left: Molecules sampled along the two interpolations. Right: Interpolation in the latent space visualized via the first two principal components. In the equivariant representation, we visualize a straight blue line between $z_1$ and $z_2$. In the invariant representation, we visualize the linear interpolation between $s(z_1)$ and $s(z_2)$ (red), and the equivariant linear interpolation between $z_1$ and $z_2$ subsequently mapped to $\mathcal{Z}_s$ (blue).

Theorems & Definitions (5)

• Proposition 4.1
• Definition 4.2
• Definition 4.3
• Proposition 4.4
• Proposition 4.5