Finding Symmetry in Neural Network Parameter Spaces

Authors: bo  zhao
nima  dehmamy
robin  walters
rose  yu
Venue: NeurIPS_cc_2024_Workshop_UniReps
Published: 2024-09-24 06:44:43

Abstract: Parameter space symmetries, or loss-invariant transformations, are important for understanding neural networks' loss landscape, training dynamics, and generalization. However, identifying the full set of these symmetries remains a challenge. In this paper, we formalize data-dependent parameter symmetries and derive their infinitesimal form, which enables an automated approach to discover symmetry across different architectures. Our framework systematically uncovers parameter symmetries, including previously unknown ones. We also prove that symmetries in smaller subnetworks can extend to larger networks, allowing the discovery of symmetries in small architectures to generalize to more complex models.

Main paper: 
 introduction: 
 Parameter space symmetry, or loss-invariant transformation of parameters, influences various aspects of deep learning theory. Continuous symmetry connects groups to their orbits, revealing important topological properties such as the dimension [Zhao et al., 2023b] and connectedness [Zhao et al., 2023a] of the minimum. Parameter symmetry also influences training dynamics through the associated conserved quantities of gradient flow [Kunin et al., 2021] and by steering stochastic gradient descent towards certain favored solutions [Ziyin, 2024] . Additionally, symmetry provides a tool to perform optimization within a loss level set, with successful applications in accelerating optimization [Armenta et al., 2023 , Zhao et al., 2022] and improving generalization [Zhao et al., 2024] . Other applications of parameter space symmetry include model compression [Ganev et al., 2022 , Sourek et al., 2021] and reducing the search space for efficient sampling in Bayesian neural networks [Wiese et al., 2023] .Despite the wide range of applications, our knowledge of parameter space symmetries is limited. In particular, known symmetries often cannot account for all loss-invariant parameter transformation. While several frameworks have been developed to unify known symmetries, whether the symmetries in current literature are complete remains an open question. Due to a lack of systematic approach, current practice typically requires deriving symmetries from scratch for every new architecture, creating barriers for wider application that leverages parameter symmetries. In this paper, we discuss an automated approach to directly learn the symmetry groups and their group actions on the parameter space of neural networks. We show that large networks often have symmetries inherited from its components or subnetworks. This view suggests that searching for symmetries in small networks is an effective approach to identify a significant number of symmetries in modern architectures.Our main contributions are:• Formal definitions of data-dependent parameter symmetries and their infinitesimal form.• An approach to identify symmetries in the parameter space of large networks from known symmetries in smaller subnetworks. • A framework that discovers symmetry in neural network parameter spaces. 
data-dependent group action and symmetry: 
 Let Θ be the space of parameters and D be the space of data. In this paper, we consider loss functions of the form L : Θ × D → R, which map parameters and a single data point to a real number. By abuse of notation, we allow L to simultaneously process multiple data points. Specifically, we sometimes defineL : Θ × D d → R d for d ∈ Z + data points.Let G be a group. Consider a map a, which defines a map for every data batch of size d ∈ Z + :a : D d → (G × Θ → Θ) X → (a X : g, θ → θ ′ ).(1)The map a is a group action on Θ if it satisfies the following axioms:identity: a X (I, θ) = θ, ∀X ∈ D d , ∀θ ∈ Θ. associative law: a X (g 2 , a X (g 1 , θ)) = a X (g 2 g 1 , θ), ∀g 1 , g 2 ∈ G, ∀X ∈ D d , ∀θ ∈ Θ.A group action a is a parameter space symmetry of L if it additionally satisfies loss invariance:L(a X (g, θ), X) = L(θ, X), ∀g ∈ G, ∀X ∈ D d , ∀θ ∈ Θ.A function L has a G-symmetry if there exists a loss-invariant group action a. We refer to G as a symmetry group of L. Additionally, the action a is termed a data-dependent group action or symmetry if the map (1) has a non-trivial dependency on X. That is, a is data-dependent if there exists X 1 , X 2 ∈ D d , such that a X1 ̸ = a X2 . We derive an infinitesimal version of parameter space symmetries in Appendix B. 
building symmetries from known ones: 
 One way to identify symmetries in a large network is by examining its components or subnetworks. Despite often having billions of parameters, neural networks typically consist of a limited set of functional families, such as fully connected layers, attention mechanisms, and activation functions. This modular view suggests a mechanism by which symmetries in networks with fewer layers might extend to those in deeper networks. Additionally, within similar types of networks, it may be possible to extrapolate symmetries found in narrower layers to wider ones.By focusing on symmetries in small architectures and using them to infer symmetries in larger ones, we circumvent the complexity associated with direct handling of high-dimensional parameter spaces. This approach not only simplifies the discovery of symmetries in large-scale networks but also provides a systematic method for using symmetries in smaller subnetworks to understand those in more extensive architectures. Proofs and further discussions can be found in Appendix C.When a loss function L depends on a subset of parameter exclusively through a subnetwork f , any symmetries that preserves f will also preserve the original network L: Proposition 3.1. Let L : Θ × D d → R d be a function, where the parameter space Θ is a product space Θ = Θ 1 × Θ 2 , with spaces Θ 1 , Θ 2 . Suppose there exist functions h :Θ 1 × D d → S, f : Θ 2 × S → T , and j : (Θ 1 × T ) × D d → R d , such that for every θ = (θ 1 , θ 2 ) ∈ Θ and X ∈ D d , L(θ, X) = j (θ 1 , f θ 2 , h(θ 1 , X) ), X . If a : S → (G × Θ 2 → Θ 2 ) is a G- symmetry of f , then there is an induced G-symmetry of L, a ′ : D d → (G × Θ → Θ), defined by a ′ X (g, (θ 1 , θ 2 )) = θ 1 , a h(θ1,X) (g, θ 2 ) .The relationship between the functions in the proposition is described by the commutative diagram below, where p 1 : Θ → Θ 1 , p 2 : Θ → Θ 2 are projections onto Θ 1 and Θ 2 , i 1 : Θ 1 → Θ 1 and i 2 : Θ 2 → Θ 2 are identity maps, and X ∈ D d represents a batch of data. When L can be decomposed in this way, the function h does not depend on Θ 2 , and the function j depends on Θ 2 only through the output of f . This effectively confines L's dependency on Θ 2 to the transformation defined by f , ensuring that any transformation on Θ 2 not altering the output of f will not affect the output of L. Consequently, symmetries identified in the smaller network f can be extrapolated to the larger network L.Θ R d Θ 1 × Θ 2 × Θ 1 Θ 1 × Θ 2 × S Θ 1 × T L(•, X) p 1 × p 2 × p 1 i 1 × i 2 × h(•, X) i 1 × f (•, •) j(•, X)Proposition 3.1 can be applied to construct symmetries in larger networks from those in smaller ones (Corollary C.2, C.1 in Appendix C). Figure 1 shows the subset of parameters (Θ 2 ) the symmetry applies to in the corollaries. Note that this approach does not explore the emergence of new, more complex symmetries that may arise as the neural network scale up in size. Notably, there are cases where there exists a G symmetry over its input space, but group actions on individual layers are not loss-invariant (Kvinge et al. [2022] ). Nevertheless, studying smaller and simpler networks remains a effective strategy to obtain a significant number of symmetries in larger networks, and is a first step in characterizing the complete set of symmetries in modern architectures.( &# # = $ & ((# # ) ( &# # = ((# # ) (a) (b) (c) ℎ , ℎ ! ∘ + ℎ ! , = ℎ ! ∘ -! Θ 0 Θ ! -" Θ " 1 " 1 ! Θ " Θ ! Θ " Θ ! 
automatic discovery of parameter symmetries: 
 Formulating symmetries in the infinitesimal form makes them easier to learn using an automatic framework, as it defines a set of local conditions for a function to be a symmetry. Using the infinitesimal symmetry derived in Section B.1, we construct an automated framework for discovering parameter space symmetries.Enforcing Loss Invariance and Group Axioms. Given a function L, our goal is to find a symmetry a and a set of Lie algebra elements h corresponding to a symmetry group of L. We parameterize a using a neural network with learnable parameters, and set h to be learnable as well. We define the following loss terms that quantify the deviation from loss invariance and the group axioms (identity and associativity law):L invariance = E x,θ |D θ L| θ,X • D g a X | I,θ (h)| (2) L id = E x,θ ∥a x (I, θ) -θ∥ 2 (3)The two loss terms bias the action towards being loss-invariant and preserving identity. By minimizing L Lie_deriv , we ensure that the learned symmetry a and the Lie algebra element h satisfy the infinitesimal symmetry condition (Theorem B.1). Minimizing L id enforces the identity axiom. By focusing on the Lie algebras, we enforce the loss invariance and group structure at the infinitesimal level. This formulation allows us to avoid computing exponential maps.Regularizations. To prevent the group action to be the identity function, we encourage the infinitesimal action to be nonzero. In implementation, we include the following regularization term to encourage the norm of the infinitesimal action to be around a fixed positive real number β:L reg_id = min a,h E θ |β -∥D g a X | I,θ (h)∥|.When learning multiple generators simultaneously, we want them to be orthogonal. Following Yang et al. [2023b] , we do this by including the following cosine similarity between each pair of the k generators in the loss function: L reg_h_orth = 1≤i<j≤k hi•hj ∥hi∥∥hj ∥ . Finally, we encourage sparsity of h for easier interpretation, with L reg_h_sparse = k,j |h kj |.The final training objective is a weighted average of the above loss and regularization terms, with hyperparameters γ 1 , ..., γ 6 ∈ R + : min h,a In the first set of tasks, we see if our method can learn generators for architectures with already known data-independent symmetries. We consider two-layer networks in the form ofEQUATIONL(W 1 , W 2 , X, Y ) = ∥W 2 σ(W 1 X) -Y ∥ 2 , where W 2 ∈ R m×h , W 1 ∈ R h×n are parameters, X ∈ R n×k , Y ∈ R m×kare data, and σ is a homogeneous activation function.During training, we train the generators h and the group action a under objective (4). We parametrize a using a 4-layer MLP with hidden dimensions 64, 64, 64. The group aciton a takes a group element, parameter, and data as input and outputs transformed parameters. We use 10000 training samples, each containing a randomly generated set of parameters and data. We set the learning rate as 10 -3 with decay 0.6 every 1000 steps, and the weights for the multi-objective loss as γ 1 = 10, γ 2 = γ 4 = γ 5 = 1, and γ 6 = 0.1.As a proof of concept, we training a group action and a single generator h ∈ R 2×2 for the two-layer architecture with m = h = n = k = 1 and σ being the identity function. Figure 2 visualizes the learned generator, which matches the expected generator that generates the rescaling group.Note that, however, we do not impose constraints on the group action (in particular, not enforcing linear actions). Hence we do not expect the learned generators to look similar to the elements of the Lie algebra infinitesimal generators of the symmetry group in general. For example, the action a can be a composition of two function, the first transforming learned generators to the set of actual generators, and the second performing the group action. We find that our method can learn the generators and group actions for wider two-layer homogeneous architectures as well. More examples of learned generators for larger architectures can be found in Appendix D.As a more practical application of our framework, we attempt to uncover data-dependent symmetries from architectures where no continuous symmetry is known before. We apply our framework to learn generators and loss-invariant group actions for two-layer neural network with sigmoid and tanh activation function, as well as a three-layer neural network with skip connection.Specifically, we aim to learn symmetries in the two-layer networks defined in the previous section, but replacing σ by sigmoid or tanh. Our objective is again to find a set of generators h and a group action a that minimizes (4). We use 10000 training samples, each containing a randomly generated set of parameters and data. We set the learning rate as 10 -3 and the weights for the multi-objective loss as γ 1 = 1, γ 2 = γ 4 = 10, γ 5 = 1, and γ 6 = 0.1.Figure 3 shows the learned generators for data-dependent symmetries in a two-layer sigmoid MLP with parameters dimensions 6 in the Appendix shows the training curve. Since sigmoid networks have no data-independent continuous symmetry, this set of symmetries are data-dependent, indicating that our method successfully learns data-dependent symmetries for this architecture.W 1 ∈ R 3×3 , W 2 ∈ R 3×1 and data X ∈ R 3×1 , Y ∈ R 1×1 . FigureFigure 4 shows the learned generators for data-dependent symmetries in a three-layer tanh MLP with parameters dimensionsW 1 ∈ R 2×2 , W 2 ∈ R 2×2 , W 3 ∈ R 2×1 and data X ∈ R 1×2 , Y ∈ R 1×1 .The generators indicate the existence of symmetries that act on non-contiguous layers, which has not been discovered in previous literature. ∈ R 2×2 , W 2 ∈ R 2×2 , W 3 ∈ R 2×1 and data X ∈ R 1×2 , Y ∈ R 1×1 . 
discussion: 
 While our discovery framework suggests that there are previously unknown data-dependent symmetries in various neural network architectures, the existence and number of symmetries in neural network parameter spaces remain open questions. Whether the number of symmetries is affected by existence of symmetry in data or changes during training are also interesting directions. Future work will examine the structure of learned symmetry, such as the dimension of Lie algebras.Proof sketch. Consider a smooth curve γ(t) = a X (exp(ht), θ) in Θ, where h ∈ g and t ∈ R. Then, since L is invariant under a, L(γ(t), X) = L(θ, X), ∀t ∈ R. The result follows from differentiating both sides with respect to t at t = 0 and applying chain rules.Proof. Since a is a symmetry of L, we haveL(a X (g, θ), X) = L(θ, X), ∀g ∈ G, ∀θ ∈ Θ, ∀X ∈ D d .Consider a smooth curve γ(t) = a X (exp(ht), θ) in Θ, where h ∈ g and t ∈ R. Then, since L is invariant under a, L(γ(t), X) = L(θ, X), ∀t ∈ R.Differentiating both sides with respect to t at t = 0, we getd dt L(γ(t), X) t=0 = 0.Applying the chain rule,d dt L(γ(t), X) t=0 = D θ L| θ,X dγ(t) dt t=0 . Now, compute dγ(t) dt t=0using the chain rule:dγ(t) dt t=0 = d dt a X (exp(ht), θ) t=0 = D g a X | I,θ d dt exp(ht) t=0 .Since exp is the exponential map from gl(n) to GL(n), and h ∈ gl(n), we haved dt exp(ht) t=0 = h. Therefore, dγ(t) dt t=0 = D g a X | I,θ (h).Putting it all together, D θ L| θ,X (D g a X | I,θ (h)) = 0.Equation 5 states that the gradient of the loss function L with respect to the parameters θ is orthogonal to the directions in parameter space generated by the infinitesimal symmetry transformations D g a X I,θ (h). This orthogonality implies that moving along these symmetric directions does not change the loss to first order, reflecting the invariance of L under the group action.Assuming that Θ = R n , then for a single data point (d = 1), we can write (5) in coordinates asEQUATIONWhen Θ = R n and G is a subgroup of GL(n) with a linear, data-independent symmetry a x (g, θ) = gθ for all x ∈ X, (6) reduces to the equation in Theorem 3.1 in [Moskalev et al., 2022] . With (D g a) ijk = ∂ai ∂g jk = δ ij θ k , we haveEQUATIONOur symmetry acts on parameters instead of data, but otherwise this matches Theorem 3.1 in [Moskalev et al., 2022] .A more straightforward training objective exponentiates the Lie algebra to obtain group elements, before enforcing loss invariance and group axioms:min h,a L invariance_int + L id_int + L assoc_int with L invariance_int = E x,θ,t ∥L (a x (exp(ht), θ), x) -L(θ, x)∥ L id_int = E x,θ ∥a x (I, θ) -θ∥ L assoc_int = h1,h2∈g E x,θ a exp(h1)X (exp(h 2 ), a X (exp(h 1 ), θ)) -a X (exp(h 2 ) exp(h 1 ), θ) .Similarly to the infinitesimal version, this objective also directly enforces the necessary group structures. We adopt the infinitesimal formulation to avoid the computational overhead of evaluating exponential maps.This section contains the proofs for results in Section 3.Proposition 3.1. Let L : Θ × D d → R d be a function, where the parameter space Θ is a product space Θ = Θ 1 × Θ 2 , with spaces Θ 1 , Θ 2 . Suppose there exist functions h :Θ 1 × D d → S, f : Θ 2 × S → T , and j : (Θ 1 × T ) × D d → R d , such that for every θ = (θ 1 , θ 2 ) ∈ Θ and X ∈ D d , L(θ, X) = j (θ 1 , f θ 2 , h(θ 1 , X) ), X . If a : S → (G × Θ 2 → Θ 2 ) is a G- symmetry of f , then there is an induced G-symmetry of L, a ′ : D d → (G × Θ → Θ), defined by a ′ X (g, (θ 1 , θ 2 )) = θ 1 , a h(θ1,X) (g, θ 2 ) .Proof. We need to show that a ′ satisfies the identity and associative law of a group action and preserves L.Since a is a group action on Θ 2 , it satisfies the identity axiom a h(θ1,X) (I, θ 2 ) = θ 2 . Applying this in the definition of a ′ , we get a ′ X (I, (θ 1 , θ 2 )) = (θ 1 , a h(θ1,X) (I, θ 2 )) = (θ 1 , θ 2 ). Since a is a group action on Θ 2 , it satisfies the associative law a h(θ1,X) (g 2 g 1 , θ 2 ) = a h(θ1,X) (g 2 , a h(θ1,X) (g 1 , θ 2 )), for all g 1 , g 2 ∈ G. It follows that a ′ also satisfies the associative law: a′ X (g 2 g 1 , (θ 1 , θ 2 )) = (θ 1 , a h(θ1,X) (g 2 g 1 , θ 2 )) = (θ 1 , a h(θ1,X) (g 2 , a h(θ1,X) (g 1 , θ 2 ))) = a ′ X (g 2 , a ′ X (g 1 , (θ 1 , θ 2 ))) Finally, since a is a symmetry of f , we have f (a h(θ1,X) (g, θ 2 ), h(θ 1 , X)) = f (θ 2 , h(θ 1 , X)), for all g ∈ G.It follows that a ′ preserves the value ofL: L(a ′ X (g, θ), X) = j (θ 1 , f a h(θ1,X) (g, θ 2 ), h(θ 1 , X) ), X = j (θ 1 , f θ 2 , h(θ 1 , X) ), X = L(θ, X).The first corollary describes how symmetries identified in narrower networks also apply to wider networks. A function σ : R h×k → R h×k is row-wise if, for any matrix A ∈ R h×k with rows {a i ∈ R k } h i=1 , the output matrix σ(A) has rows {σ row (a i ) ∈ R k } h i=1 , where σ row : R k → R k applies independently on each row of A. Element-wise function is a special case of row-wise functions. For fully connected networks with row-wise activation functions, identifying a symmetry in one architecture suggests that the same symmetry will apply to wider versions of that architecture.Corollary C.1. Consider a network parameter space Θ(m, h, n) = R m×h × R h×n and data space D(n, k) = R n×k . Let σ : R h×k → R h×k be a row-wise function. Consider a function L mnhk : Θ(m, h, n) × D(n, k) → R m×k , defined as L mnhk ((U, V ), X) = U σ(V X) for U ∈ R m×h , V ∈ R h×n , and X ∈ R n×k . If there is a G-symmetry of L mnhk , then there is a G-symmetry of L mnh ′ k with any h ′ > h.In addition to obtaining symmetries from those in smaller networks, we can also get symmetries for a loss function over data batches with a certain size, if we know there is a symmetry for this function over larger data batches. Concretely, if there exists a group action that preserves loss for all data batches of size d ∈ Z + , then that group action preserves loss for all data batches of size d ′ < d. Preprint. 
 
 Appendix:  


Decision: In light of the positive reviewers' feedback and relevancy of the submission, we are pleased to accept this paper for presentation at UniReps 2024. We kindly ask the authors to incorporate the reviewers' suggestions and feedback in the final camera-ready version of the manuscript.

Reviews: review: title: title: Review of Symmetry Discovery in Neural Network Parameter Spaces
main_review: main_review: ## Summary
The paper presents a novel attempt to discover parameter space symmetries in neural networks systematically. It does so by first introducing data-dependent parameter symmetries and deriving their infinitesimal form. demonstrating that symmetries in smaller subnetworks can extend to larger networks. The authors propose an automated approach to uncover symmetries in neural network architectures, including previously unknown ones by imposing loss invariance and group axioms such as identity and associativity, while incorporating regularization to prevent trivial solutions and ensure meaningful symmetry discovery.

## Strengths
- The framework automates the discovery of symmetries in neural network parameter space, reducing the manual effort required to identify such transformations, and extending the existing knowledge of parameter space symmetries.
- By showing that symmetries in small subnetworks can apply to larger architectures, the framework potentially simplifies the analysis of complex models.
- The paper is written in a clear and engaging style, which makes it enjoyable to read.

## Weaknesses
- While symmetries are extended from small to larger networks, the potential for emergent symmetries in larger architectures is not thoroughly investigated.
- Current validation focuses on smaller architectures. The scalability of the framework to large, real-world architectures is not fully explored, and further empirical validation is needed.
- Although the paper is well-written and engaging, certain sections contain notations that can be challenging to follow. The mathematical density in parts may pose difficulties for readers not deeply familiar with the subject matter. Providing a clearer exposition or introducing the notations earlier would further enhance the accessibility of the content.

### Minor Comments:
There are two grammatical errors in lines 30 and 125: "identifies" and "training" to be taken to their base forms.

## Evaluation
- The theoretical contributions are solid and well-supported by empirical validation on small architectures. The formal framework for symmetry discovery is novel and could open up new directions in neural network research.
- The idea of systematically discovering symmetries in neural network parameter spaces, particularly data-dependent symmetries, is highly original. The automated approach to learning these symmetries represents a significant advancement.
- The potential applications of this work in neural network training, optimization, and generalization are substantial. This paper addresses a gap in current literature, where most known symmetries are derived manually.
title: title: Review
main_review: main_review: This paper presents a framework for identifying and analyzing parameter space symmetries in neural networks. The authors formalize data-dependent parameter symmetries and develop a method to uncover symmetries in smaller subnetworks that can generalize to larger models. The authors conduct several preliminary experiments to demonstrate their approach. Specifically,  they show that  the framework is able  to discover both known data-independent symmetries in simple MLP architectures and uncover previously unknown data-dependent symmetries in more complex networks.


Strengths:
- The method seems novel to me. The presented idea looks interesting.
- The framework seems general and could be used to explore additional architectures.

Weaknesses:
- The current results are only shown on small architectures, but I think that for a preliminary workshop it is fine.
- It is not clear what is the computational cost of the method.
- What are the assumptions on the data? How much data do we need to learn symmetries?
title: title: Official Review of Submission60 by Reviewer Bs1w
main_review: main_review: **Summary:**  

The paper presents a framework for discovering data-dependent symmetries in neural networks by utilizing infinitesimal symmetries and Lie group generators. It efficiently extends symmetries identified in smaller components to larger architectures. The method reveals both known and novel symmetries, providing valuable insights into the underlying structure of neural networks.

**Strengths and Weaknesses:**

**Strengths:**  

- The authors offer a solid theoretical justification for analyzing the hidden symmetries in smaller network components to infer the overall symmetries of the full architecture. This approach is valuable because it enables researchers to study the properties of individual components and still contribute to a broader understanding of weight space symmetries.
- As demonstrated in (Zhao et al., 2023b), data-dependent non-linear action symmetries can play a crucial role in uncovering hidden symmetries. These actions enable researchers to explore weight space symmetries beyond traditional intertwiner groups linked to activation functions (Godfrey, Charles, et al., 2022). However, designing non-trivial non-linear actions is challenging. The authors address this by proposing a method that learns group actions while discovering the corresponding group's Lie algebra generators.

**Weaknesses:**  

- In line 96, $L_\text{Lie deriv}$ is mentioned without prior definition. A clear explanation of this term should be provided.
- The derivatives used in Equation 2 should be explicitly defined before introducing the loss function. A formal definition, similar to the style used in Theorem B.1, would enhance clarity.
- I see two key challenges in learning both the Lie algebra generators and non-linear actions simultaneously:
  1. The dimensionality and number of generators are difficult hyperparameters to tune. 
  2. As the authors acknowledge, the learned action function can preprocess the generators into an intermediate state before mapping them to the final group action, reducing the interpretability of the learned symmetry.

**Questions:**  
1. Do you have any heuristics for selecting the dimensionality and number of generators?
2. Once the generators and corresponding group actions are learned, one possible application could involve training a new equivariant or invariant network that uses the weights of the previous networks as input, leveraging the learned symmetries (Tran, Hoang V., et al., 2024). While methods like frame averaging (Puny, Omri, et al., 2021) could be applied in this scenario, do you have any alternative approaches or potential applications in mind?

**Limitations:** Limitations are adequately addressed

**Ethical concerns:** No

**Soundness:** 3 good

**Presentation:** 2 fair

**Contribution:** 4 excellent
title: title: Good paper
main_review: main_review: After presenting a framework of symmetry discovery and decomposing the symmetry of subnetworks, this paper proposes to optimize over (1) the group action approximated by a neural network, (2) the infinitesimal generator of the parameter symmetry group, so that the loss is invariant.
From my limited knowledge and understanding of this paper, I have several questions and comments:

1. Contributions in Sections 2 and 3 are somewhat related but orthogonal. The formulation in Section 2 correctly formalizes the motivation but is not practically significant, nor directly used in the rest of the paper.
2. The group action $a:(g,\theta)\mapsto\theta'$ is a very complicated 4-layer neural network. Will the network memorize the parameter so that it does not depend (much) on $g$? If this is the case it may hurt the usefulness of $g$ and I would suggest constraining the form of the action to be simpler by analyzing the form of the neural network.
3. The hyperparameter tuning seems very rough to me, which reduces the practical usefulness. It may lead to the vanishing of certain terms. At least, could they be tuned to equalize every loss term? For example in Figure (2), will the block-diagonal structure be more visible if the orthogonality loss term dominates?
4. I don't understand why it suffices to calculate the first-order condition at $g=I$ only.

The paper is overall well presented and theoretically sound. Good paper.