Paper: Latent Space Symmetry Discovery (ICLR-2024)

Reviewer comment (Theory):
As discussed in the paper, the nonlinear-symmetric-discovery problem itself is ill-posed, and there are many meaningless "optimal" solutions to it due to the representation power of neural networks. This paper incorporates several patches to alleviate this issue, such as an orthogonal weight matrix in the final layer, and zero-mean of the latent features within an empirical batch. These regularization terms seem strong and hard-coded, and there is no theoretical understanding/analysis of them.

Author response:
* The regularizations introduced in Sec 4.3 are necessary to address the failure modes in latent symmetry discovery. They are not hard-coded but well-motivated by the experiments. We provide more detailed discussion below.

Reviewer comment (Evaluation):
Are there metrics distinguishing the qualities of learned symmetries other than human interpretation?

Author response:
* **Yes**. A possible quantitative metric is the *equivariance error*: $E_{x,g}\|gf(x)-f(gx)\|^2$. Also, as Reviewer 5mQG has suggested, the *logit invariance* proposed by Moskalev et al (2023) are also related metrics. We have included these metrics in **Appendix A.6**.

Reviewer comment (Theory):
Can the model, after applying all these regularization terms, learn all desired symmetries?

Author response:
* **Yes**, these regularizations won’t remove any symmetries. First, the orthogonal parametrization would not cause any problem, because if we have an encoder with non-orthogonal final layer $W$, we can always do the Gram-Schmidt process to get an orthogonal weight $Q=PW$, which is effectively a change of basis in the latent space. Whatever linear symmetries in the original latent space should remain, possibly with a different group representation. Thus, in theory, an orthogonal final layer suffices to learn all desired symmetries. The zero-mean normalization would not affect most symmetries but does remove translation symmetry. However, the current group representation $v \mapsto \pi(g)v$ where $\pi(g)$ is a matrix cannot express translation of $v$ anyway, so the zero-mean normalization is compatible. To discover translation symmetries, though, we need to expand the search space to the affine group: $V \rtimes GL(V)$ and remove the normalization.

Reviewer comment (Theory):
Are these regularization terms enough to rule out all meaningless solutions?

Author response:
* We find that these regularizations resolve some common failure modes in the current experiments, but it is very difficult to enumerate all the unseen failures and rule them all out.

Reviewer comment (Experiments):
Similarly, all learned/discovered symmetries in the experiments are rotation-based. Why is that? Is it related to the choice of the regularization? Can the model learn other symmetries in practice? For example, can it learn a nonlinear version of E(n)?

Author response:
* The success of symmetry discovery relies on proper choices for the hyperparameters, such as the dimensions of the Lie algebra and its representation, and the weights for the reconstruction loss and the GAN loss. Here are some principles for hyperparameter tuning. The principle for choosing the Lie algebra dimensionality is to start from a small value and gradually increase it until the generator cannot fool the discriminator, because it is required to produce more symmetries than there actually exist. Then, we choose a latent dimension that leads to a reasonably small reconstruction loss. This ensures that there is no information loss and the encoder and the decoder can be inverses of each other. The loss weights depend on different tasks, because the scales of the reconstruction loss vary.
* **Yes, we provide some additional results of non-rotation-based symmetry and in domains other than synthetic dynamical systems.** The first example is the scaling symmetry in Lotka-Volterra system (**Appendix A.3**).
Another example is **Top Tagging**, which is also used in LieGAN (Yang et al 2023). The task is a binary classification between top quark jets and the background signals. The original dataset is invariant to Lorentz transformations. To have a nonlinear symmetry, we apply a nonlinear transformation to the input: we choose 4 spatial modes $u_1,u_2,u_3,u_4 \in \mathbb R^{128}$ given by Legendre polynomials and define $x_{\text{new}}=\sum_{i=1}^4 x_iu_i$, where $x=(x_1,...,x_4)$ is the 4-momentum from the original dataset. (Note that this is really just an arbitrary choice of nonlinear transformation that does not have much physical meaning. Just a quick procedure to remove the linear symmetry and show our method’s ability to discover nonlinear group actions.) The discovered Lie algebra is [visualized HERE](https://postimg.cc/1gqmRN8m). Its representation is different from Figure 5 from Yang et al (2023), because the latent reps obtained by the encoder are different from original inputs. However, we can compute the structure constants of this Lie algebra ([visualized HERE](https://postimg.cc/gXbYxkhv)), which reveal its similar structure to the Lorentz algebra $\mathfrak{so}(1,3)$.

Reviewer comment (Evaluation):
The learned nonlinear symmetries are not that interpretable due to the neural encoder.

Author response:
* The lack of interpretability comes naturally with the task of learning nonlinear group actions. It is often difficult to arrive at simple, closed-form formulas that describe all the nonlinear transformations considered, especially for high-dimensional systems. In our method, we managed to decouple the nonlinearity into neural components that do not rely on groups and group elements. Thus, we can analyze the group structure from the learned linear representations. This provides some degree of interpretability on the types of learned symmetries (e.g. rotation, scaling).

Reviewer comment (Experiments):
How difficult is it to learn a meaningful nonlinear symmetry using this method in practice? Are there results showing this method learned symmetries other than those rotation-based? Are there results in domains other than the synthetic dynamic systems?

Author response:
* The success of symmetry discovery relies on proper choices for the hyperparameters, such as the dimensions of the Lie algebra and its representation, and the weights for the reconstruction loss and the GAN loss. Here are some principles for hyperparameter tuning. The principle for choosing the Lie algebra dimensionality is to start from a small value and gradually increase it until the generator cannot fool the discriminator, because it is required to produce more symmetries than there actually exist. Then, we choose a latent dimension that leads to a reasonably small reconstruction loss. This ensures that there is no information loss and the encoder and the decoder can be inverses of each other. The loss weights depend on different tasks, because the scales of the reconstruction loss vary.
* **Yes, we provide some additional results of non-rotation-based symmetry and in domains other than synthetic dynamical systems.** The first example is the scaling symmetry in Lotka-Volterra system (**Appendix A.3**).
Another example is **Top Tagging**, which is also used in LieGAN (Yang et al 2023). The task is a binary classification between top quark jets and the background signals. The original dataset is invariant to Lorentz transformations. To have a nonlinear symmetry, we apply a nonlinear transformation to the input: we choose 4 spatial modes $u_1,u_2,u_3,u_4 \in \mathbb R^{128}$ given by Legendre polynomials and define $x_{\text{new}}=\sum_{i=1}^4 x_iu_i$, where $x=(x_1,...,x_4)$ is the 4-momentum from the original dataset. (Note that this is really just an arbitrary choice of nonlinear transformation that does not have much physical meaning. Just a quick procedure to remove the linear symmetry and show our method’s ability to discover nonlinear group actions.) The discovered Lie algebra is [visualized HERE](https://postimg.cc/1gqmRN8m). Its representation is different from Figure 5 from Yang et al (2023), because the latent reps obtained by the encoder are different from original inputs. However, we can compute the structure constants of this Lie algebra ([visualized HERE](https://postimg.cc/gXbYxkhv)), which reveal its similar structure to the Lorentz algebra $\mathfrak{so}(1,3)$.

Reviewer comment (Evaluation):
Are there more ways to interpret/use the learned symmetries?

Author response:
* **Yes**. Just as equivariant networks work on inputs with linear symmetries, after our method learns a symmetry with a linear action on the latent space, we can pass the latent embeddings to subsequent equivariant layers for downstream tasks. For example, we can make better use of the learned symmetry in the equation discovery algorithm by **enforcing the equivariance in SINDy regression**. Here is the idea in brief. SINDy learns a linear combination of basis functions: $\dot z=f(z)=WD(z)$, where $W$ is the learnable coefficients and $D$ is the function space basis. Mathematically we can show that if there is equivariance $g\dot z=f(gz)$, where $g$ is in a Lie group with $L$ in its Lie algebra, then $LWD(z)=W\frac{\partial D}{\partial z}Lz$. Similar to Equivariant MLP (Finzi et al, 2021), we can solve this linear constraint wrt trainable parameter $W$ and write $W=Q\beta$, where $Q$ is a constraint matrix and the vector $\beta$ becomes the actual trainable parameter. In this way, we are effectively using the learned symmetry as a constraint and compressing the search space of equations.

Reviewer comment (Theory):
There is a strong emphasis on the non-linear group action aspect in this paper, but I believe this is a bit misguided. This is because the hallmark result of representation theory of Lie groups is that the Lie algebra connects the group to the vector space. Moreover, this can be described by matrices---hence linear representation---and you do not generally need non-linear representations. In practice, however, you can codify non-linear actions (e.g. rotations of 3D objects in a 2D image) and this is where you might want to learn a non-linear action. But I find the emphasis on the non-linear action exaggerated because LaLieGan learns a linear rep in the latent space anyways. I would suggest toning down these claims.

Author response:
* We believe that discovering nonlinear group actions is an important aspect. Although representation theory is built on linear representations, in some scenarios, no linear action can describe the symmetries of the systems. This is illustrated in all our dynamical system experiments where LieGAN (Yang et al, 2023) fails. You mentioned that Lie algebra elements can be described by matrices, which is indeed an important fact. However, it is important to distinguish the group elements and their transformations on data space. The linear reps can describe the Lie algebra (and Lie group) elements, but they are not enough to describe their transformations on an arbitrary data manifold.
* Therefore, we build the nonlinear action around a linear rep to enable symmetry discovery while also remaining compatible with the representation theory. As pointed out by the other reviewers, nonlinearity is an important next-step solution compared to existing methods e.g. LieGG (Moskalev et al, 2022) and LieGAN (Yang et al, 2023).

Reviewer comment (Theory):
Prop 4.1 seems to not apply to the setup that the authors consider. This is because the encoder and decoder map to a latent space of an autoencoder. This means that the latent dimension can be **lower** than the observation dimension. As a result, $\phi$ and $\psi$ cannot ever be inverses---i.e. bijective---because the information is lost. Thus, I have strong doubts about the value of the proposition. Moreover, many symmetry discovery methods already assume an autoencoder setup. The main difference is that they do not take an adversarial approach so this limits the novelty of the method. Finally, the paper learns approximate inverses anyways so there is no reason to guarantee that the learned representation is an exact Lie group.

Author response:
* In fact, Prop 4.1 **can** apply to our setup. In many tasks (e.g. the reaction-diffusion system in our paper), it is possible to encode the high-dimensional observations into a lower-dimensional latent space without loss of information. This is because the high-dimensional observations dwell on a low-dimensional manifold, and the intrinsic dimensions of these systems can be much lower. In such cases, the encoder $\phi$ is the inverse of the decoder $\psi$ restricted to the data manifold $\mathcal M \subset V$. Formally, Prop 4.1 is indeed problematic because it does not state this restriction to the data manifold. It should modified to: if $\phi|_{\mathcal M}$ and $\psi$ are inverse of each other, where $\mathcal M$ is the data manifold, then $\pi’(g, \cdot)$ is a valid group action. Thanks for pointing this out!
* Regarding the novelty, our method does not require a priori knowledge of the symmetry group or additional information such as group element associated with each data point or interaction with environment. The latter is closely related to your next question and we provide further discussions there.
* Finally, it is true that Prop 4.1 relies on an exact inverse condition. However, even when this condition is slightly violated in practice, it is still possible to learn a good group action. In Appendix C.2, we evaluate the validity of the learned group actions on the test dataset via the reconstruction loss and the compatibility error.

Reviewer comment (Theory):
One of my biggest concerns is that the approach and results in this paper go against a relatively known result in Linear Symmetry Based Disentanglement by (Caselles-Dupré et. al 2019) who prove that symmetry discovery is impossible without interaction with the environment. This result is a symmetry-based analog to the result by Locatello et. al 2019. Thus I fear that the results in this paper are generally not true, and going beyond the toy datasets considered here might be impossible.

Author response:
**No, our method does not contradict the results in SBDRL (Caselles-Dupré et al 2019),** because the task settings and assumptions are different:
* By symmetry, we refer to the equivariance of a function $f$, i.e. $f(gx)=gf(x)$. For example, $f$ can be the evolution function in a dynamical system: $x_{t+1}=f(x_t)$. But it would be more helpful to think of $f$ as an arbitrary function, e.g. an image classifier / a time series forecaster, written as $y=f(x)$. Then, our dataset $\\{(x,y)\\}$ consists of input-output pairs of this function. We discover the equivariance of the function from the input-output pairs.
* In Caselles-Dupré et al (2019) and Quessard et al (2020), the transitions from $o_t$ to $o_{t+1}$ rely on a group element $g$ instead. Their datasets are trajectories of $(o_0,g_0,o_1,g_1,...)$. They learn a map $f$ from observation $o\in W$ to latent $z\in Z$ that is equivariant between group actions on $W$ and $Z$.

It should be noted that **many other works in symmetry discovery** belong to the first category. They **do not rely on interactions with the environment either.** (Zhou et al, 2021; Desai et al, 2022; Dehmamy et al, 2021; Yang et al, 2023)

Despite the differences, SBDRL (Caselles-Dupré et al 2019), as well as its theoretical results, is surely related to our work. Next, we provide some discussions about the connection. For this part, we will follow the notations in the SBDRL paper.

Theorem 1 in SBDRL states that there are multiple worlds equipped with different group actions on the state space $W$ that can produce the same dataset. Thus, the symmetry-based representation $(f, \cdot_Z)$ in one world is necessarily not equivariant in another world. This is indeed true, but it only means that there are multiple possible representations of symmetry, and our method can learn one of them.

Using the hue change example (on page 4 immediately below Thm 1) from this paper, the succession of colors obtained by the right translation (i.e. $C_3$ group action on the state space) can be $(r, b, g, r)$ or $(r, g, b, r)$. For the first one, we can define equivariant representation as $f: r \mapsto 0, b \mapsto 1, g \mapsto 2$, and $hz = z + h\ \mathrm{mod}\ 3$ for $h \in \\{ 0, 1, 2 \\} = C_3$. This is indeed not a symmetry-based representation of the other world, but if we define the group action on the latent space as $hz = z - h\ \mathrm{mod}\ 3$ it becomes a valid representation.

The setting of Thm 1 assumes there is a privileged group action attached to the environment and it is important to learn exactly that one. Thus it is problematic to have multiple $(W_i,\cdot_{W_i})$ which agree with training data and some of which are inconsistent with the learned $(f, \cdot_Z)$.

In our case, however, we just wish to learn a group action which helps to understand the structure of the data manifold and can be used in downstream applications to reduce the dimensionality of the problem. Thus, each $(W_i,\cdot_{W_i})$ could be realized by a different $(f_i, \cdot_{Z_i})$, and we are satisfied with any of them. Thus, we believe that there is no contradiction.

In retrospect, it might have been useful to specify the degree of non-uniqueness in the specification of our learning problem, i.e. given $(f, \cdot_Z)$, what can be done to produce a new $(f’, \cdot’_Z)$ which is just as good? Can we derive some equivalent relation $R$ and get the quotient space $\mathcal H / R$, where $\mathcal H \ni (f, \cdot_Z)$ denotes our hypothesis space? Answering these questions can reduce the complexity of our learning problem and may be a good direction for future work.

Reviewer comment (Experiments):
I am confused as to why the authors do not compare with more standard baselines for symmetry discovery. Granted these works often assume knowledge of the group apriori but why this is not the correct test bed? For example $SO(N)$ is done in Fig 2 and 4 of Quessard et al 2020 (you even cite this paper) as well as the main experiment of Caselles-Dupré et. al 2019. Moreover, there has been a lot of development in Deep Delay auto-encoders that extend SINDy. In particular Deep Delay Autoencoders Bakkarji et. al 2023 is an appropriate baseline for the non-linear dynamical system discovery experiments. I encourage the authors to include this baseline as well.

Author response:
**We didn’t compare because the settings of Quessard and Caselles-Dupré are different from ours, as discussed previously.** Assuming knowledge of the group is a minor difference. The main difference is that they consider a trajectory of observations and group transformations: $(o_0,g_0,o_1,g_1,...)$, while we aim to discover the equivariance of a function: $gf(x)=f(gx)$.

But we can compare under a slightly modified setting, where we manually define such a function that is equivariant to cyclic translations. **We include the Flatland experiment as in Quessard and Caselles-Dupré in Appendix A.5, where we show that our method can learn similar 4D latent representations as in Figure 2 from Quessard.** 

Finally, regarding Delay SINDy autoencoder, there is no publicly available code. We tried to implement it by ourselves, but it did not yield good discovery results. There are many loss terms that are jointly optimized, and we find it difficult for all the reconstruction and consistency losses to converge simultaneously. We will not include this baseline for now. Any other suggestions on equation discovery baselines are welcome.

We hope that our response can provide clarifications toward your questions. Please let us know if you have additional arguments against the paper or the rebuttal. If not, may we kindly request you to increase the score for our paper?

Reviewer comment (Experiments):
In Eq. 2 you learn transformations as $\sum_i\text{exp}[w_i L_i]$. How to choose the number of matrices L to be used in the method? I suppose the number of Lie algebra elements you parametrize will significantly affect the flexibility of the method in the latent space.

Author response:
* The number of matrices L, i.e. the dimension of Lie algebra, is a hyperparameter. A principle for choosing this hyperparameter is to start from a small value and gradually increase it until the generator cannot fool the discriminator, because it is required to produce more symmetries than there actually exist. In this way, we can discover the maximal symmetry. In the original experiments, we only considered one-dimensional Lie algebra. We provide an example of learning a two-dimensional Lie algebra for the Lotka-Volterra system in **Appendix A.3**.

Reviewer comment (Theory):
In Proposition 4.1 you mention that $\psi$ and $\phi$ are inverse of each other. It is not correct, these functions are inverse to each other only on the input dataset. It raises the following question: how robust is the inverse property when you move away from the training dataset? How robust is the detected symmetry, when you move away from the training dataset? It reminds me of the following paper: *Moskalev A. et al. On genuine invariance learning without weight-tying. Topological, Algebraic and Geometric Learning Workshops 2023. – PMLR, 2023*.

Author response:
* Prop 4.1 relies on the inverse condition. However, even when this condition is slightly violated in practice, it is still possible to learn a good group action. We show an example in Appendix C.2, where we evaluate the robustness of the learned group actions on the test dataset via the reconstruction loss and the compatibility error.  
* The paper from Moskalev proposes some useful metrics for measuring learned invariance. We can use our discriminator on the test set to evaluate the *logit invariance*. Also, another possible metric is *equivariance error*: $E_{x,g}\|gf(x)-f(gx)\|^2$. We have included these metrics in **Appendix A.6**. Thank you for pointing us to this paper!