MatrixNet: Learning over symmetry groups using learned group representations

Authors: lucas  laird
circe  hsu
asilata  bapat
robin  walters
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:16:38

Abstract: Group theory has been used in machine learning to provide a theoretically grounded approach for incorporating known symmetry transformations in tasks from robotics to protein modeling. In these applications, equivariant neural networks use known
symmetry groups with predefined representations to learn over geometric input data. We propose MatrixNet, a neural network architecture that learns matrix representations of group element inputs instead of using predefined representations. MatrixNet achieves higher sample efficiency and generalization over several standard baselines in prediction tasks over the several finite groups and the Artin braid group. We also show that MatrixNet respects group relations allowing generalization to group elements of greater word length than in the training set. Our code is available at https://github.com/lucas-laird/MatrixNet.

Main paper: 
 introduction: 
 The choice of representation for input features is a key design aspect for deep learning. While a wide variety of data types are considered in learning tasks, for example, sequences, sets, images, time series, and graphs, neural network architectures admit only tensors as input. Various methods exist for mapping different input types to tensors such as one-hot embeddings, discretization of continuous signals, learned token embeddings of image patches or words [1, 2] , adjacency matrices [3] , positional encodings [1] , or spectral embeddings [4] .In this paper we consider the question of what feature representations to use for learning tasks with inputs coming from a symmetry group. There are many examples of tasks defined over symmetry groups, such as policy learning in robotics [5] , reinforcement learning [6] , pose estimation in computer vision [7] , sampling states of quantum systems [8] , inference over orderings of a set [9] , and grouptheoretic invariants in pure mathematics. Past work has typically employed fixed representations chosen from among the known representation theory of the group. Representation theory is the branch of mathematics concerned with classifying the set of representations of a group G which, in this context, refers to homomorphic realizations of a group in terms of n × n matrices. For groups with well understood representation theory, for example, the symmetric group S n or SO(n), this provides a ready set of embeddings for converting group elements into tensors for use in downstream models. Trial and error or topological analysis has shown that the choice of group representation is critical for learning [10, 11, 12, 13 ]. 1Instead of using predefined group element representations, we propose to learn feature representations for each group element using a learned group representation. That is, we learn to map group elements to invertible n × n matrices which respect the symmetry group structure. There are several advantages to this strategy. First, unlike predefined group representations, learned representations allow the model to adapt to the given task and capture relevant information for the learning task. Second, learned representations provide reasonable correlations between the learned features for different group elements since they incorporate algebraic structure into the model. This structure is encouraged using free group generators and group relation regularization. Third, relative to learned vector embeddings, learned matrix representations are very parameter efficient for encoding different group elements, reducing the problem to learning only representations of group generators; using sequence encoding, our method is able to generalize to combinatorially large or even infinite groups. Fourth, the learned representation admits analysis in terms of the irreducible subspaces of the generators, giving insight into the model's understanding of the task.We integrate our learned group representation into a specialized architecture, MatrixNet, adapted for learning mappings related to open problems in representation theory. We compare against several more general baselines on order prediction over finite groups and estimating sizes of categorical objects under an action of the Artin braid group. Through our experiments we observe that our approach achieves higher sample efficiency and performance than the baselines. We additionally show that MatrixNet's constraints allow for it to generalize to unseen group elements automatically without the need for additional data augmentation.Our contributions include:• Formulation of the problem of learning over groups as a sequence learning problem in terms of group generators and invariance to group axioms and relations,• The MatrixNet architecture, which utilizes learned group representations for flexible and efficient feature learning over symmetry groups,• The matrix block method for constraining the network to respect the group axioms and an additional loss term for learning group relations,• Empirical validation showing MatrixNet outperforms baseline models on a task over the symmetric group and a task over the braid group related to open problems in mathematics. 
related works: 
 Mathematically Constrained Networks Many deep learning methods incorporate mathematical constraints to better model underlying data. For instance, the application of graph neural networks [3, 14] to problems with an underlying graph structure has led to state of the art performance in many domains. Deep learning models have also been designed for tasks with known mathematical formulations by parameterizing components of algorithmic solutions as neural networks and leveraging their structures for more efficient optimization [15, 16] . More broadly the field of geometric deep learning [17] advocates for building neural networks which reflect the geometric structure of data in their latent features. When symmetries are present in data, group equivariant neural networks [18, 19, 20, 21] can enable improved generalization and data efficiency by incorporating known symmetries into the model architecture using the representation theory of groups. Our method also utilizes group representations, but unlike equivariant neural networks, we use learned as opposed to fixed representations. We also focus on modeling functions defined on the group as opposed to between representations of the group.Learning Structured Representations Instead of using predefined representations as inputs, many methods seek to learn mathematically structured representations from data. This idea has been applied in physics [22, 23] , robotics [24] , world models [25] , self-supervised learning [26, 27] , and unsupervised disentanglement [28] . Park et al. [29] , for example, use a combination of learned symmetry and equivariant constraints to map images to group elements or vectors in group representations. Similar techniques are used in symmetry discovery where the underlying group symmetry is not known a priori [30, 31, 32] . Yang et al. [32] use a generative model to learn latent representations with a linear group action in order to find unknown group symmetries in data.Here, in contrast, we start with a known group and learn a matrix representation. Hajij et al. [33] propose algebraically-informed neural networks which learn a non-linear group action from a group presentation. We also learn a group action but consider linear group actions and the learning signal comes not only from the group presentation but also a downstream task.Deep Learning for Math Recent work has shown deep learning can be useful for providing examples, insight, and proofs related to open problems in mathematics. One approach is the application of language models to mathematics [34] , which has the benefit of flexibility in how the model is prompted yet is difficult to interpret and prone to errors. Meanwhile significant work has been done in the area of symbolic regression and automated theorem proving [35] . Other work applies deep learning to the direct modeling of partial differential equations [36, 37] . These methods can perform exceptionally well on real-world data [38] , but suffer when trying to interpret predictions made for the purposes of mathematical research. Another avenue involves training graph neural networks on mathematical data such as group or knot invariants and analyzing the learned representations to see which features are significant as a way to provide intuition to mathematicians [39] . Our method also uses structured inputs and learned features, but uses a sequence model and learned group representation instead of a graph neural network with learned node attributes.3 BackgroundQuestion: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The Limitations section details the assumptions of the model and the limitations of regularization as a method for enforcing group relation invariance.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
group theory: 
 Groups encode systems of symmetry and have been used in machine learning to build invariance into neural networks to various transformations [18] . Formally, a group G is a set equipped with an associative binary operation • : G × G → G which satisfies two axioms: (1) there exists an identity element 1 ∈ G, such that g • 1 = 1 • g = g, (2) for each g ∈ G, there exists an inverseg -1 ∈ G such that g • g -1 = g -1 • g = 1.Examples of groups include finite groups such as the dihedral group D 4 which gives the symmetries of the square, SO(3), the continuous group of 3D rotations, or (Z, +) the infinite discrete group of integer shifts.Since groups may be combinatorially large or infinite, it is essential to encode their elements and compositional structure in a succinct way. For many discrete groups, generators and relations provide such a description. A set of elements S = {g 1 , ..., g n } ⊆ G are called generators if every element of G can be written as a composition of g 1 , g -1 1 , ..., g n , g -1 n . In general, each element of G may be written in many different ways in terms of the generators; this non-uniqueness is encoded using a set of relations. A set R = {r 1 , . . . , r m } of words in the generators S are relations for G if each word r i is equal to the identity in G and if R generates the entire set of words equal to identity under composition and conjugation. The generators and relations of a group taken together are called a presentation and denoted G = ⟨g 1 , ..., g n | r 1 , ..., r m ⟩. For example D 4 = ⟨r, f | r 4 = f 2 = f rf r⟩. Due to relations, group elements do not have a unique word representation. For example f rf = r 3 = r 7 all represent the same group element. By convention relations are sometimes stated as equalities instead of single group elements, for example f rf = r 3 . The free group F S = ⟨g 1 , . . . , g n ⟩ is defined to have no relations except for those coming from the two group axioms above.An important notion for our discussion is the order of an element. If g ∈ G then the order of g, denoted |g|, is the smallest k such that g k = e. (For non-finite groups, k may be infinity). The order of the group |G| is simply the number of elements in the group. For any g, Lagrange's theorem implies that |g| is a divisor of |G| [40] , which restricts the possible orders an element may take.Abstract group presentations are difficult to work with in many settings. Group representations map group elements to invertible matrices such that composition of group elements corresponds to matrix multiplication. This gives the group a natural action on vector spaces and allow for analysis of the group using linear algebra. Formally, a representation of a group G is a group homomorphismΦ : G → GL(n) to the group of invertible n × n matrices. That is Φ(g 1 • g 2 ) = Φ(g 1 ) • Φ(g 2 ).A property of Φ is that Φ(1) = I n×n and Φ(g -1 ) = Φ(g) -1 . Due to the homomorphism property, it is sufficient to define Φ for generators of the group G. For example, a 2 × 2 representation of D 4 is given by mapping r to a π/2-rotation matrix and f to a reflection over the x-axis.The representations of many groups are well classified. This provides a ready source of tensor representations for group elements to use as inputs for neural networks. For example, for finite groups, by Maschke's Theorem [41] , representations may be decomposed into irreducible representations. That is, there exists a basis such that the representation matrices are all block diagonal with the same block sizes and these blocks cannot be further subdivided. The irreducible representations may then be further classified by computing character tables.The braid group on n strands B n has presentationEQUATIONThe braid group intuitively represents all possible ways to braid a set of n strands. The generators σ i correspond to twisting strand i over i + 1 and σ -1 i is the reverse, twisting strand i + 1 over i. It is defined topologically as equivalence classes up to ambient isotopy of n non-intersecting curves in R 3 connecting two sets of n fixed points. The braid group is infinite and though some representations are known, they are not fully classified [42] . The braid group has important connections to knot theory, mathematical physics, representation theory, and category theory.The symmetric group on n elements, denoted S n , is defined as the set of bijections from {1, . . . , n} to itself. It is also a quotient of the braid group B n and has a presentation similar to Eqn. 1 but with additional relationsσ 2 i = 1 for 1 ≤ i ≤ n -1.Here σ i is the transposition (i i+1). The symmetric group has finite order |S n | = n!. Representations of the symmetric group are well understood. The irreducible representations are parameterized by partitions of n. For more details, see [43] .One current active research problem in mathematics concerns actions of braid groups on categories. A category is an abstract mathematical structure that has objects and maps or morphisms between objects, satisfying several coherence axioms. For example, the category Vect R has objects which are real vector spaces and morphisms which are linear maps. Functors are maps between categories that take objects to objects and morphisms to morphisms between the corresponding objects, satisfying several compatibility conditions. The action of a group G on a category C means that each group element g ∈ G is associated with an invertible functor F g : C → C, such that any relation x = y in the group implies that the corresponding functors F x and F y are naturally isomorphic.Given a category on which a braid group acts, mathematicians are interested in measuring how objects grow under repeated applications of elements of the braid group. The "size" of an object in a category may be measured using a tool called Jordan-Hölder filtrations. For example, Bapat et al. [44] attempt to measure growth rates of objects in a specific category C n under repeated applications of certain twist functors σ Pi , which define an action of the braid group B n on C n . Each object in the category has a Jordan-Hölder filtration giving a unique vector in Z n ≥0 of Jordan-Hölder multiplicities. For more details see [44] and Appendix A. However, they are only able to compute the action in certain cases and a simple formula is elusive. A complete description of the Jordan-Hölder multiplicities after applying combinations of σ Pi to one of the generating objects is only known for n = 3; that is, the case of the 3-strand braid group B 3 . Understanding how these multiplicities evolve under repeated application of braids is a challenging open research problem in mathematics.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: The hardware used is reported in the appendix. Our model and data do not require significant compute and are small enough to be run on most machines.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
methods: 
 We formulate the problem of learning a function on a symmetry group as a sequence learning problem using a presentation of the group in terms of generators and relations. We propose MatrixNet which predicts the label using a learned matrix representation for the group. The homomorphism property of the representation is enforced through a combination of model design and an auxiliary loss term.We consider task functions of the form f : G → R c where G is a finite or discrete group and the output space R c may represent either a regression target or class label. While such tasks appear in computer vision, robotics, and protein modeling, we are particularly interested in problems in mathematics where neural models may lend additional examples and insight towards proving theorems.To efficiently represent group elements in infinite or large groups, we consider a presentation of G = ⟨S | R⟩ in terms of generators S = {g 1 , . . . , g n } and relations R = {r 1 , . . . , r m }. Model inputs g ∈ G are represented by sequences of generators (g i1 , . . . , g i ℓ ) where g = g i1 • . . . • g i ℓ is of arbitrary length ℓ ≥ 0 and 1 ≤ i j ≤ n. For convenience, we can include the identity g 0 = e among the generators to pad sequences without changing the group element.Since a single group element may be represented by different sequences, it is critical for the model f θ to be invariant to both the group axioms and relations. That is, we desire We propose MatrixNet (see Figure 1 ), a neural sequence model which models functions on a group G by taking as input sequences of generators for G. It achieves invariance to group axioms and relations through a combination of built in constraints and an auxiliary loss term. The key part of the MatrixNet architecture is the matrix block which takes a group generator g i as input and outputs an invertible square matrix representation M gi . The matrix representation M g for an arbitrary group element g is defined as the product of matrix representations of generators needed to generate g. This matrix representation is then flattened and passed to a downstream task model (such as an MLP) which computes the label. In what follows, we give a more detailed description of the matrix block and some variations on the architecture.EQUATION)EQUATIONf θ (g i1 , . . . , g i k , r j , g i k+1 , . . . , g i ℓ ) = f θ (g i1 , . . . , g i k , g i k+1 , . . . , g i ℓ ) (Relations). (4)We define the signed one-hot encoding, a modified version of the traditional one-hot encoding, for encoding group generators used as input to MatrixNet. Let (g ϵ1 i1 , g ϵ2 i2 , ..., g ϵ ℓ i ℓ ) be a sequence of generators where 0 ≤ i k ≤ n and ϵ k ∈ {±1}. The signed one-hot encoding encodes each generator g ϵ k i k as a vector v gi k = ϵe i = [0, . . . , 0, ϵ k , 0, . . . , 0] ∈ R n which is 1 or -1 in the ith entry. The identity element g 0 = 1 is mapped to the zero vector v 1 = 0 ∈ R n . The signed one-hot encoding is chosen since it intuitively relates a generator and its inverse as v g -1 = -v g .The matrix block is designed as a parameterized representation of the free group F S , that is, a homomorphism Φ : F S → GL(n) which satisfies 3 properties: (1) the matrixΦ(g) = M g is invertible, (2) Φ(1) = I n×n , and (3) if Φ(g) = M g then Φ(g -1 ) = M -1 g .In what follows v gi k is the signed one-hot encoding of the generator g i k and W is a learnable parameter matrix. For a group element g = g i1 . . . g in , the matrix block is definedA k = Reshape(W v gi k ) M gi k = MatrixExp(A k ) M g =M gi 1 M gi 2 . . . M gi nwhere Reshape reshapes a vector in R n 2 into a square n × n matrix. Proposition 1. Matrix Block defines a representation of the free group.Proof. Property ( 1) is satisfied since the outputs of a matrix exponential are invertible. Properties ( 2) and ( 3)follow from v g -1 i k = -v gi k , v 1 = 0,and properties of the matrix exponential,M 1 = MatrixExp(0 n×n ) = I n×n M g -1 i k = MatrixExp(-A k ) = M g -1 i k .We now present some variations of this simple design that satisfy the group homomorphism properties.Linear Network (MatrixNet-LN) The first variant replaces the single parameter matrix W with a linear network that has two parameter matrices W 1 , W 2 . The linear network matrix block changes the computation of the intermediate A k matrix to:A k = Reshape(W 2 W 1 v gi k )This is still a linear function of v gi k meaning Proposition 1 still holds by the same argument. While this variation does not give increased expressivity over the original formulation of the matrix block, the linear network can change the optimization landscape leading to different performance in practice. This variation is called MatrixNet-LN.Non-Linearity (MatrixNet-NL) The second variation introduces an element-wise odd non-linear function f . That is f (-x) = -f (x) as with tanh. The non-linear matrix block modifiesA k = Reshape(f (W 1 v gi k ))Proposition 2. Non-linear matrix block defines a representation of the free group.Proof. Property (1) is satisfied by the same argument in Proposition 1. Since f is an odd function, f (W 1 v 1 ) = f (0) and Reshape(f (W 1 v g -1 i k )) = Reshape(f (-W 1 v gi k )) = -A k . ThereforeA k . That is, A k = Reshape(W 2 f (W 1 v gi k )). Unless otherwise noted we set f = tanh.The third variation is inspired by the decomposition of representations into irreducible representations. For certain classes of groups such as finite groups, every representation decomposes such that the matrices M g have a consistent block diagonal structure in some basis. Thus to learn an arbitrary representation of the group G, it suffices to learn a block diagonal representation assuming the blocks are large enough. This variation learns ℓ intermediate n j × n j matrices A kj which are combined to form a blockdiagonal n × n matrix A k where n = ℓ j=1 n j . The block-diagonal matrix block is defined withA kj = Reshape(W j v gi k ) for j = 1 to ℓ A k = BlockDiag(A k1 , A k2 , . . . , A k ℓ )Note that MatrixExp and matrix multiplication preserve the block structure. If the sizes n j are fixed to all be equal, this formulation can be implemented as a multi-channel matrix block where both A k and M gi k are ℓ × n j × n j tensors with ℓ channels. Each A kj is calculated identically to A k in the original matrix block formulation, and BlockDiag is linear, so Proposition 1 still holds. This variation is also compatible with the previous two variations. In our experiments, we implement a 3-block version called MatrixNet-MC with a single linear layer and no non-linearity.The matrix block is constrained to learn a representation of the free group F S . As a consequence MatrixNet will satisfy ( 2) and (3) as desired. However, most groups have relations which cannot be enforced through a simple weight-sharing scheme used in equivariant architectures [19] . We propose to learn the relations through a secondary loss which measures how closely the representation respects the group relations. More concretely, let G = ⟨S|R⟩ be a group with relations r i ∈ R. The loss is:L rel = ri∈R (||M ri -I n×n ||)Since MatrixNet learns a representation that is invariant to group axioms, it is sufficient to sum over only {r i } m i=1 and not all compositions of relations. For architectures which do not respect the free group structure, the relations r i alone may not guarantee that all equivalent words have identical feature representations, requiring potentially combinatorial amounts of data augmentation. This allows MatrixNet to both efficiently learn group relation invariance and simply verify this invariance without any data augmentation. 
experiments: 
 We use two learning tasks to evaluate the four variants of MatrixNet and compare our approach against several baseline models. We use several finite groups on a well understood task as an initial test to validate our approach and then move on to an infinite group, the braid group B 3 , on a task related to open problems. As baselines, we compare to an MLP for fixed maximum sequence length and LSTM and Transformer models on longer sequences. Baseline model parameters were chosen so all of the models have approximately the same number of trainable parameters.The first task is to predict the order of group elements in finite groups. Elements of finite groups are input as finite sequences of generators as described in Section 3.3. The typical efficient algorithm for computing the order involves disjoint cycle decomposition, making order classification a non-trivial task. See Appendix B.1 for more details on the sampling method and data splits.We compare MatrixNet variants and three baselines for order prediction in S 10 :(1) MLP with 3 layers with hidden dimension 256 and SiLU activations, (2) Fixed representation input to a 2-layer classifier MLP with 256 hidden dimensions and SiLU activations, (3) LSTM input to a 2-layer LSTM with 256 hidden dimensions with a subsequent MLP classifier using SiLU activations, (4) MatrixNet-LN with a 2-layer 256 hidden dimension matrix block and classifier network with SiLU activations. (5) MatrixNet-MC with a 2x2 matrix block size over 5 matrix channels and classifier network with SiLU activations. (6) MatrixNet-NL with a 2-layer 256 hidden dimension matrix block with SiLU activations and classifier network with SiLU activations. The precomputed representation is an ablated version of MatrixNet using a fixed representation of S 10 instead of a learned one. For the fixed representation, we use the standard 10×10 representation given by the permutation matrices corresponding to the group element. In MatrixNet-LN, the activation between layers of the matrix block is set to a linear passthrough while in MatrixNet-NL the activation is specified to be SiLU. MatrixNet-MC enforces a 2x2 block diagonal structure on the learned representations corresponding to the 2-dimensional irreps of S 10 .We also note the use of SiLU activation in our S 10 MatrixNet model. Due to the generator self-inverse property we need not consider separate generator inverses, and so the odd function requirement given in Proposition 2 is not applicable. MatrixNet achieves a high classification accuracy across all additional groups tested. However, accuracy for the Abelian group is lower than the accuracies for other groups tested (87% vs 99%).One explanation for this decrease in accuracy is due to the large number of valid orders of the group. Additionally, due to the structure of finite Abelian groups, many element orders will be underrepresented in random sampling.In our second experiment, we train models to predict the Jordan-Hölder multiplicities from braid words in the braid group B 3 (see Section 3.4). The task is formulated as a regression task with a mean-squared error (MSE) loss function. The Jordan-Hölder multiplicities are integers, so we evaluate accuracy by rounding the vector entries to the nearest integer. This accuracy is reported as an average accuracy over the three entries of the Jordan-Hölder multiplicities vector. Elements of B 3 are generated by two generators σ 1 , σ 2 and their inverses and are encoded using a signed one-hot encoding. For more details on the dataset generation process and data splits see Appendix B.2.We additionally performed an experiment to evaluate how well MatrixNet generalizes to unseen braid words longer than those seen in training. For this experiment, we compare against the MLP and LSTM since these were the highest performing baselines.We trained all of the models for 100 epochs as all of the models except the Transformer converged within 100 epochs. Despite performance converging much faster than 100 epochs for most MatrixNet variants, we found that additional epochs of training improved the model's invariance to group relations with minimal variations in performance. M σ2 M σ1 = M σ2 M σ1 M σ2 ||.We also compute this distance between two non-equivalent braids σ 1 σ 1 σ 2 , σ 2 σ 2 σ 1 for reference under Non-Relational Difference in Table 4 .The relational error results in Table 4 mirrors the extrapolation performance confirming that representation quality is important for effective generalization. High relational error compounds over longer word lengths hindering generalization whereas low relational error allows MatrixNet to automatically generalize to longer word lengths through invariance to group relations. These results show that MatrixNet, particularly the MatrixNet-LN and MatrixNet-NL variants, is able to learn group representations invariant to the group relations allowing for effective generalization to longer unseen group words. We present some visualizations of the learned representations of the braid group from the highest performing variant, MatrixNet-NL. Figure 2 shows visual plots of the learned representations. In the last two plots, the learned representation for two equivalent words are approximately equal even though this relation is not among the relations r j used in the loss L rel . This shows MatrixNet does indeed generalize over relations, allowing it to generate nearly identical representations for equivalent words.  
conclusion: 
 In this paper we have presented MatrixNet, a novel neural network architecture designed for learning tasks with group element inputs. We developed 3 variations of our approach which structurally enforce group axioms and a regularization approach for enforcing invariance to relations. We evaluate MatrixNet on two group learning problems over several finite groups and B 3 and demonstrate our model's performance and ability to automatically generalize to unseen group elements. In future work we plan to develop interpretability analysis methods based on group representation theory to better understand the structure of MatrixNet's learned representations. Understanding the learned representations may provide valuable insights and explanations of the model outputs assisting with generating new conjectures for open mathematical research problems.Limitations The current work relies on the assumption that the studied group is finitely presented which limits us to discrete groups. However, learned group representations may also be useful for learning over Lie groups. In such case, extending our method will require working with infinitesimal Lie algebra generators. Additionally, while the group axioms are strictly enforced by the model structure, the fact the relations are enforced using auxiliary loss terms means the homomorphism property is not exact. Future work may explore methods of reducing this error.Recall that the path algebra of a directed graph (or quiver) Q over some field k is generated as a free vector space by all possible paths in Q, including the trivial paths at each vertex. The product in kQ is given as follows: let q : a → b be a path and p : c → d be a path. The product pq is equal to zero unless b = c. If b = c, then the product pq is simply the composite path a q -→ b p -→ d. The path algebra of Q is denoted kQ.We are interested in a quotient of the path algebra of Γ dbl n called the zig-zag algebra, and denoted Z n . To obtain Z n , we impose the following relations on kΓ dbl n . In what follows, (i|i ± 1) represents the unique arrow i → i ± 1.• For each i, set(i + 1|i)(i + 2|i + 1) = 0 and (i -1|i)(i -2|i -1) = 0. • For each i, set (i + 1|i)(i|i + 1) = (i -1|i)(i|i -1).Consequently, in the zig-zag algebra, any path of length at least 3 is automatically zero, and the only surviving paths of length 2 are back-and-forth loops starting from any vertex (and all possible such loops are set to be equal). The paths of length 0 and 1 remain as-is.Let Z nproj be the category of (graded) projective modules over Z n . The category C n is constructed as a differential graded version of the bounded homotopy category of complexes of projective modules over Z n , in which we identify the internal grading shift with the homological grading shift, and consequently also the triangulated shift. For more details, see, e.g. [44, Section 6] . 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This work focuses entirely on mathematical problems with no broader societal impacts.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not use any pretrained models and generate our own data. The data is not sensitive or private in nature.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We cite all packages used and acknowledge compute resources used during this research. We do not use any preexisting assets in this research.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: We are not releasing new assets as part of this submission. We will provide documentation and code upon publication. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In this section, we record some important properties of the category C n . First, the category C n is triangulated: there is a a triangulated shift functor [1] : The category C n is generated as a triangulated category by the objects P 1 , . . . , P n . That is, the smallest triangulated subcategory of C n (closed under isomorphisms) that contains the objects P 1 , . . . , P n is C n itself. These objects correspond to the indecomposable projective modules of the zig-zag algebra Z n .C n → C n which is an equivalence. An n-fold composition of [1] is denoted [n].Each object P i is spherical. This means thatHom n (P i , P i ) = k n = 0 or n = 2, 0 otherwise.Remark A.1. The reason for the notation is that the ring of endomorphisms of P i of all possible degrees is isomorphic to the cohomology ring of a sphere (in this case a 2-sphere).It is a general fact that any spherical object X of a triangulated category gives an associated functor σ X , called the spherical twist in X (see Seidel-Thomas [45] for more details.) The functor σ X is an auto-equivalence; that is, it has an inverse equivalence σ -1 X such that compositions in both directions are isomorphic to the identity functor.In particular, we obtain equivalences σ Pi : C n → C n .The category C n has a bounded t-structure. This means that there is an abelian subcategory A n ⊂ C n , such that every object X ∈ C n has a unique finite filtration whose factors lie in A[i] for decreasing values of i. This filtration is called the cohomology filtration. In fact, this abelian subcategory A n is also generated by the objects P i : it is the extension-closure of the objects P i . Moreover, the objects P i are simple objects of A n .First consider any object X ∈ A n . The category A n is a finite-length abelian category. It is a standard fact that X has a Jordan-Hölder filtration whose factors are simple objects in A n , namely the objects P i . Now consider a general object X ∈ C n . We first consider the cohomology filtration of X, with factors Y j ∈ A[j]. For each Y j ∈ A[j], we consider its (shifted) Jordan-Hölder filtration, which breaks Y j up further into copies of P i [j]. Putting these two together, we obtain a finer filtration of the object X, which we also call the Jordan-Hölder filtration of X.The factors of this Jordan-Hölder filtration are shifted copies of P i for all i. Thus we can count the number of occurrences of each P i in the Jordan-Hölder filtration of X, and it is well-known that these counts do not depend on the specific choice of the Jordan-Hölder filtration.Recall the spherical twist functors σ Pi : C n → C n . A remarkable observation of Seidel-Thomas [45] is that these functors (weakly) obey the relations of the n-strand braid group. That is, we have isomorphisms of functorsσ Pi σ Pj ∼ = σ Pj σ Pi whenever |i -j| ̸ = 1, and σ Pi σ Pj σ Pi ∼ = σ Pj σ Pi σ Pj whenever |i -j| = 1.Since these are precisely the relations of the group B n , we obtain an action of B n on the objects of the category C n .Consider the following broad question. Given an object X of C n and a braid g ∈ B n , can we relate the Jordan-Hölder multiplicities of β(X) to the Jordan-Hölder multiplicities of X? Stated more simply, can we compute the Jordan-Hölder multiplicities of β(P i ) for any i and any β?Answers to this question are known in some cases, and remain open in others. For instance, a complete answer was obtained by Rouquier-Zimmermann [46] for the 3-strand braid group B 3 acting on C 3 . This answer was rediscovered and refined in terms of more general filtrations (Harder-Narasimhan filtrations) in [44] and [47] .It is also known that if β = σ ℓ Pj for some ℓ, then the limit as ℓ → ∞ of the counts of β(X) can be obtained, up to a common scaling factor, by computing the sum of the dimensions of Hom m (P j , X) for all m [47] .However, for the vast majority of values of n and most of the elements of the braid groups B n , we do not have a good answer to this question. While there is vast potential for future work, we write down a few specific open problems.1. Generalise the Rouquier-Zimmermann theorem (and its corresponding versions in [44] and [47] ) to larger values of n. 2. We can compute a finer version of Jordan-Hölder multiplicities: split up the number of occurrences of each P i by degree shift. That is, record the number of occurrences of P i [d] separately for every possible d. This information can be encoded in a polynomial in one variable in q ±1 , in which the coefficient of q d is the multiplicity ofP i [d].Generalise the Rouquier-Zimmermann theorem in this setting to larger values of n. 3. By using a more refined version of Jordan-Hölder multiplicities, known as Harder-Narasimhan multiplicities, we observe that the possible Harder-Narasimhan factors of any object of the form β(P i ) are highly constrained, and satisfy some very nice combinatorial properties. This constraint can be explicitly described for any B n via a geometric model (due to Khovanov-Seidel [48] ) for objects in the category C n . Nevertheless, the relationship of these constrained sets with the action of B n is mysterious. For example, given an object X, is there an algorithm to write down a braid that will send X to an object with a desired set of Harder-Narasimhan filtration factors?Length Extrapolation Experiment For the generalization experiment we generated a dataset of all braid words up to length 7 which was split into 80% training and 20% validation. We also generated three separate test datasets of 10,000 examples each with braid words of length 8, 9, and 10 to evaluate how performance degrades over increasing length. The models were trained for 100 epochs on the training and validation sets and then tested on the three test sets.Regularization Details All MatrixNet architectures were trained using the regularization loss defined in Section 4.3. We chose the Frobenius norm for the norm and used the braid relationσ 1 σ 2 σ 2 = σ 2 σ 1 σ2 and the inverse σ -1 1 σ -1 2 σ -1 1 = σ -1 2 σ -1 1 σ -1 2 .The regularization term was added to the loss every 10 training batches.Hardware Details All of the categorical braid action experiments were run on a machine with a single Nvidia RTX 2080 ti GPU.We also performed a length interpolation generalization experiment to test how well each model generalizes to braid words that are shorter than the maximum length seen during training. The results are not presented in the main body of the paper as all models generalize to shorter braid words. Justification: We provide bulleted claims and contributions in the introduction. We provide experimental results and proofs in the methods and experiments sections.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Justification: This paper focuses on deep learning in a mathematical context. There were no human participants or personal data involved in this research.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).38th Conference on Neural Information Processing Systems (NeurIPS 2024).The term representation is used in both deep learning and mathematics with related but different meanings. We will disambiguate the former as a feature representation and the latter as a group representation. 


Decision: The paper proposes to learn symmetry groups so that, in turn, to learn optimal representations for downstream tasks at hand. In particular, MatrixNet is proposed, which takes as input generators, then combines them in a Matrix Block that returns an invertible square matrix. Then for a new group element, its matrix representation is the product of the matrix representations of the respective generators that are needed to generate the group element. A parameterized neural network form of the Matrix Block is proposed, as well as variations. An interesting experiment is on a mathematical problem of predicing Jordan-Holder multiplicities, which to be frank is out of my domain, but seems relevant and interesting for the community.

The reviews are mixed, with most reviewers appreciating the novelty and the presentation of the paper. The main complaint is whether the experiments represent a real-world setting, and if a mathematical problem constitutes really a real-world problem. I think it does, or at least it is an interesting new direction to consider for possible ways to validate machine learning models.

Further, while reading the paper, I think some subsections of related work are missing, specifically on 'Symmetry Discovery', eg 

Yang et al., Latent Space Symmetry Discovery, PMLR, 2024
Gabel et al., Learning Lie Group Symmetry Transformations with Neural Networks, TAGML, 2023
Dehmamy et al., Automatic Symmetry Discovery with Lie Algebra Convolutional Networks, NeurIPS, 2021

and in 'Mathematically Constrained Networks' a subsection on mathematical programming for imbuing constraints to neural networks:

Pervez et al., Differentiable Mathematical Programming for Object-Centric Representation Learning, ICLR 2023
Pervez et al., Mechanistic Neural Networks for Scientific Machine Learning, ICML 2024

I suggest the reviewers incorporate the reviewers' suggestions and extend their related work accordingly.

Reviews: review: paper_summary: paper_summary: The authors use neural networks to learn group representations. A group element is represented by its generators formatted as a sequence of learned matrix representations. These generators are mapped to a single matrix representation of the group element via the Matrix Block which enforces group axioms. The resulting feature is used for downstream tasks on groups. The authors show that the proposed architecture MatrixNet successfully predicts group element orders for S_10 and Jordan-Holder multiplicities for braid group B_3. Further analysis on word length extrapolation and visualization demonstrates the superiority and usefulness of the approach.
strength_weakness: strength_weakness: 1. It’s unclear why naive MatrixNet and MatrixNet-MC cannot extrapolate to longer word length. 
2. Empirical evaluations are a bit limited. It might be helpful to evaluate groups with different properties (see questions).
questions: questions: In the matrix block, is it possible to introduce commutative operations between different matrix representations of group generators if the input group element comes from an abelian group. It seems possible to introduce further architectural constraints for specific groups.
limitations: limitations: None
paper_summary: paper_summary: This paper studies feature representations of a group element for supervised learning. It considers a regression task where an input is a group element g of a finite group and a target is some label. Firstly, g is decomposed into a sequence of generators (g_1, ..., g_n) such that g = g_1 \circ ... \circ g_n. Next, each generator is mapped by a trainable matrix embedding W: Gen(G) --> R^{n \times n} followed by the matrix exponential so that we can get a group representation M_i = exp(W(g_i)) ∈ GL(n). Then their product M_1 \circ ... \circ M_n becomes the feature. Similar variants are also proposed. The performance is evaluated on two synthetic tasks: the order prediction of the symmetric group S_10 and the action prediction of the Braid group B_3.
strength_weakness: strength_weakness: 1. Limited applicability to real tasks. To the best of my knowledge, regression (or classification) of group elements is practically important for continuous groups such as SO(3) (e.g., pose estimation). However, the current approach is only applicable to finite groups. Also, I think the tasks conducted in the experiments are not directly bridged to real problems, and I feel uncertain about how the proposed method is practically valuable. 
1. It is unclear which component significantly contributed to the final performance gain. The proposed method consists of (at least) two parts: the decomposition of g into generators and a trainable representation of a generator. The current experiments are not designed to evaluate them separately. 
1. The experiments are not convincing enough—they are relatively small scale, use synthetic tasks only, and have less variety (two tasks).
questions: questions: 1. Given g, is its decomposition into the generators always uniquely determined?
1. In the experiment of 5.1, what will happen when we employ the group decomposition while using the fixed representation? I mean, the feature of g is given as concat[\rho(g_1), ..., \rho(g_n)] where g = g_1 \circ ... \circ g_n and \rho(g_i) is some representation of g_i (e.g. irreducible rep).

Typo: 
* One of M_{g_ik^-1} would be {M_{g_ik}}^{-1} in the equation below line 203.
limitations: limitations: Limitations are adressed.
paper_summary: paper_summary: The paper describes MatrixNet, a method to learn group representations such that they are optimized for a certain task of interest using a neural network. , The neural network takes in a group element in the form of a sequence of generators that compose to form the group element and forms an intermediate output that contain the matrices that are the group representations for each of the generators. The generators can them be multiplied to form the group representation which is then fed to further neural layers that are task-specific. The network is trained so as to (approximately) satisfy the constraints of group representations using architectural constraints and the loss function. Experiments on predicting the order of a element of the symmetric group and Jordan-Holder multiplicities for the 3-strand braid group.
strength_weakness: strength_weakness: 1. One of the weaknesses is that the parts about the braid group are very difficult to follow and probably need a lot more background for readers and attendees for this conference. This includes myself and I was not able to fully follow the details or why that experiment is important.

2. The experiments are also a little weak in my opinion. It was not clear to me how the learned representations were different from precomputed ones and why they were better for a given task. Also, what happens when the size of the group representations is larger. Perhaps for the first experiment, it would be nice to see the results as a function of the size of the group representation. I am not sure if the proposed architecture and learning mechanism can learn good representations when the size is large. This is important to address, in my opinion.
questions: questions: No additional questions.
limitations: limitations: Yes, authors have listed the limitations of the current submission and suggested directions for future research.
paper_summary: paper_summary: In this paper, authors study the question of what feature representations to use for learning tasks with inputs coming from a symmetry group. They propose MatrixNet, a neural network architecture that learns matrix representations of group element inputs instead of using predefined representations.

The main contributions are as follows:
1. Formulation of the problem of learning over groups as a sequence learning problem in terms of group generators with invariance to group axioms and relations.
2. The proposed neural network architecture, MatrixNet, achieves higher sample efficiency and generalization over several more general baselines in prediction tasks over the symmetric group and Artin braid group.
3. The matrix block method for constraining the network to respect the group axioms and an additional loss term for learning group relations.
strength_weakness: strength_weakness: The motivation is not so clear to me. The applications of group theory are everywhere in the real-world. Thus, how to choose a good representation for the group which is associated with the learning task is more important. In this paper, authors formulate this problem but do not give a clear answer (see my questions below), and authors focus on solving a more abstract group theory question instead which is already studied a lot by mathematician. 

I didn't see any insight behind the paper. For the experiment, I think it only solved a mathematical problem at simple setting which is far from what we expect. For example, how the model works when the group's order become larger?
questions: questions: 1. The motivation is to solve a mathematical problem or real-world learning task? I hope it can help me understand the paper clearly and inprove the writing.

2. In subsection 4.1, authors formulate the problem. In line 166, can $f$ be seen as a representation of $G$ or not? If it is a representation, can you write a explicite expression as an example? Since if it is a representation, the target domain $R^{c}$ may not be a linear space which makes it unclear for me.

3. Can you show more experiment results? Such as larger symmetric group? Also I didn't understand the meaning of predict the order of the group element? I think we should determine it instead of estimating it? Maybe I did't get your point, but I am open to discuss.
limitations: limitations: see the weaknesses and questions.

Comments: title: To authors
Thank you for the response
comment: Thank you for your response. My concerns are almost addressed. Still, I keep the original score because I'm not convinced enough of its applicability. Since I'm not a mathematician and I cannot judge how the group problems are significant.
Thank you for answering my questions. 

I believe the authors have done a good job in addressing most of the concerns that all the reviewers had.

The authors should definitely include all the new results in the rebuttal into the main paper.

I will raise my score to a 7.
Thank you for your comment - To test Matrixnet performance on broader classes of groups we've generated data for the following groups: $S_{12}$, $S_{5} \times S_{5} \times S_{5}\times S_{5}$, and $C_{11} \times C_{12} \times C_{13} \times C_{14} \times C_{15}$, and trained using Matrixnet-tanh with minimal tuning. We hope that the inclusion of the $S_5$ product group and latter Abelian group help illustrate the robustness of our method. Results are summarized in the table below:


| Group| Rep Size | Loss | Test Acc |
|  :----------------  | :-: | :------:  |  :----:  |
| $S_{12}$ | 12 | 1.1e-2 | 98.4%
| $S_{5} \times S_{5} \times S_{5}\times S_{5}$ | 20 | 2.1e-2 | 98.6%
| $C_{11} \times C_{12} \times C_{13} \times C_{14} \times C_{15}$ | 10 | 1.01e-5 | 100%
Thank you for the additional experiments and clarifications on how to incorporate other group constraints. My concerns are addressed. I maintain my score. 

Reasons for not raising my scores: I am not familiar with abstract algebra enough to see the full implication of this work. As of now the proposed approach works well on tasks chosen in the paper. It's hard for me to see if the specific architectural constraints here generalize to broader classes of groups.
Thanks for authors' reply. It helps me understand the paper better.