Monomial Matrix Group Equivariant Neural Functional Networks

Summary

The paper explores the important field of learning over weight spaces, where neural networks process other neural networks. Previous research has highlighted the importance of designing equivariant architectures that account for the symmetries of the input neural network, with a primary focus on permutation symmetries. However, previous literature did not account for all symmetries of the input NNs, particularly the weight scaling symmetries of ReLU networks and the weight sign flipping symmetries of sin or Tanh networks. This paper naturally extends previous approaches to account for these activation-based symmetries. This paper first formalizes the group of symmetries that includes both neuron permutations and scaling or sign-flipping transformations using monomial matrices. Next, it proposes a novel architecture for weight space networks that are equivariant to groups of monomial matrices. The new architecture is more efficient in terms of the number of trainable parameters compared to baseline weight space networks (which are only permutation equivariant).

Strengths

1. The paper is mostly well-written and well-structured. 2. The paper addresses the important and timely problem of learning in deep weight spaces, presenting a novel architecture that extends previous permutation equivariant networks to also account for scale/sign-flipping symmetries. 3. The proposed architecture is parameter efficient compared to baseline weight space networks.

Weaknesses

My main concern is the limited empirical evaluation and missing natural baselines. Also, the presented empirical results, mostly show marginal to no improvement over the limited baseline methods evaluated. The insufficient empirical study of this paper significantly damages the paper’s contribution. Given the current state of the learning over weight spaces literature, I would expect a more diverse, challenging, and comprehensive empirical evaluation. 1. The main text provides very few details on the construction of G-equivariant layers. I suggest the authors to provide at least one concrete example for mapping between a subspace of U, for example, mapping between some bias $b_i$ to a bias $b’_j$. 2. Since the method is built over NFN, it is limited in the sense that each monomial-NFN can process a specific input architecture. Building over or extending the work to GNN-based weight space networks will allow [1,2] the processing of diverse input architectures. 3. Missing weight-space baselines, like GNN-based models [1,2], DWSNets [3] (which, while mathematically equiv. to NFNs, obtains better empirical performance, see e.g., [6]) and NFT [4]. 4. Another important missing natural baseline is to use a permutation equivariant baseline like DWS/NFN or the GNN-based models together with scaling/sign-flipping data augmentations as in [5]. 5. Some evaluation and comparison of runtime and memory consumption w.r.t. baselines would be beneficial. 6. Also, adding some ablation regarding design choices would be beneficial. References [1] Graph neural networks for learning equivariant representations of neural networks, ICLR 2024. [2] Graph metanetworks for processing diverse neural architectures, ICLR 2024. [3] Equivariant Architectures for Learning in Deep Weight Spaces, ICML 2023. [4] Neural Functional Transformers, NeurIPS 2024. [5] Improved Generalization of Weight Space Networks via Augmentations, ICML 2024. [6] Neural Processing of Tri-Plane Hybrid Neural Fields, ICLR 2024.

Questions

1. Is it possible to use the proposed method together with normalization layers? 2. How easy would it be to extend the method to other activation functions?

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

Yes.

Summary

The present manuscript concerns the design of Neural Networks capable of processing the weights and biases of other Neural Networks, particularly Fully Connected and Convolutional NNs (known in the literature as *weight space networks, neural functional networks or metanetworks*). Previous works considered only hidden-neuron-permutation symmetries of weights/biases that preserve the function of the NN. In contrast, the authors highlight the existence of symmetries induced by certain activation functions (ReLU – positive scaling symmetries and sine/tanh – sign symmetries – Eq. (20)), following the works of Godfrey et al., NeurIPS’22 and Wood and Shawe-Taylor, Discr. Appl. Math.’96, where in certain cases are mentioned to be maximal (i.e. the only symmetries that preserve the NN function) as proved in the past. Combining these symmetries with permutations leads to the so-called “monomial-matrix groups”. Based on this background, the authors design monomial-matrix equivariant NNs, following the classical NN design strategy: linear layers interleaved with non-linearities. To characterise the former, they identify the weight-sharing pattern that equivariance imposes by solving a system of weight constraints for both the symmetries above (thus fully characterising linear layers), while for the former they use the same non-linearities as the activation functions at hand. Finally, they propose a method for monomial-matrix *invariant* layer design, which is combined with equivariant layers to yield an end-to-end invariant NN. The method is experimentally tested in various tasks: CNN generalisation prediction and Implicit Neural Representation (INR) classification and editing, showing competitive performance against permutation-symmetry-only baselines and a significant reduction in the number of learnable parameters.

Strengths

**Significance/Impact.** The topic of NN processing has been steadily gaining traction in the last year and has the potential to provide significant advantages in various applications such as meta-learning and processing signals of arbitrary nature (encoded into INRs) under a unifying framework. Therefore, improving computational efficiency and incorporating new inductive biases, as done in this work, is an important step towards advancing and popularising the field. **Presentation**. The paper is in general well-presented, with appropriately chosen notations and clear descriptions of the background concepts involved and the innovations proposed. **Novelty**. - *Studied problem*. The symmetries discussed in this work, although mentioned in the literature, have not been approached so far in the context of weight space networks. - *Methodology*: This work introduces novel layers that are equivariant or invariant to a group that has been underexplored (monomial matrices). - *Theory*. Additionally, following traditional weight-sharing proof strategies, the layers are characterized as the only linear layers that have the equivariance property, while Remark 4.5. borrows results from relevant papers (that might be not widely known to the community) to highlight the cases where the studied symmetries are the maximally function-preserving ones (however some results are missing – see weaknesses). **Quality/Execution**. The paper is well-motivated, provides a comprehensive background discussion, follows a rigorous and well-established methodology to design monomial-matrix group equivariant/invariant layers and provides adequate experimental comparisons, including in regimes where previous works fail. **Computational efficiency**. The proposed method leads to a significant reduction in the number of parameters, a property which is particularly useful for large NNs (a typical use case in modern deep learning).

Weaknesses

**Limited expressivity (possibly reflected in the experimental results?)** - Although characterising the linear layers by solving the weight sharing constraints is a fairly general and rigorous technique for equivariant layer design, I have the impression that the resulting weight sharing, in this case, is severe. I.e. due to the large size of the group considered, the resulting layers seem weak in terms of expressivity. For example, in Eq. (22) all hidden layer weights and bias updates are calculated by linearly processing each element individually, i.e. a weight corresponding to an edge between two neurons will be updated based only on its previous value, ignoring other edges across the same or other layers. - Additionally, the activation functions that can be used are quite limited to preserve equivariance (the authors mention that they use the same activation as in the NN that is being processed). - Although these design choices are necessary for the current construction, this probably will not be the case for a different one, i.e. by directly designing non-linear layers. In other words, I am concerned that working with the standard NN pattern (linear layers interleaved with non-linearities) might be too limiting for this family of symmetries. - I am also wondering if limitations in expressivity are induced by the choice of the invariant layers. Could the authors elaborate on this? - I do not consider the above as grounds for rejection, since I believe that even incorporating these symmetries into NFNs and the linear layer characterisation are sufficient contributions. However, they seem like important limitations, which I suspect are reflected in the experimental results, and therefore should be highlighted by the authors. **Related work and existing results**. - Given that the field is relatively new, I think that the authors should devote more space to a more detailed literature review, e.g. describing the weight symmetries that have been discovered more thoroughly and discussing/comparing the weight space networks that have been proposed so far. Currently, most methods are cited, but not adequately explained. I understand that space might not allow this, but at least adding an extended related work section in the appendix would help. - Additionally, the following very related works are missing: (1) In the topic of weight space networks: - Universal Neural Functionals, Zhou et al., arxiv’24 - Graph metanetworks for processing diverse neural architectures, Lim et al., ICLR’23 - (2) In the topic of weight space symmetries (conditions for maximality): - Reverse-engineering deep ReLU networks, Rolnick et al., ICML’20 - Hidden Symmetries of ReLU Networks, Grigsby et al., ICML’23 - As far as I understand, the works of Chen et al., Neur. Comp.’93, Fefferman et al., NIPS’93 characterise maximal weight space symmetries only for the tanh activation and not for sine, as the authors mention L208 – in case this is true, the statement in Remark 4.5. should be rectified. **Extensibility**. Judging by the derivation of the weight-sharing patterns, it appears that it is not straightforward to extend this approach to symmetries induced by other activation functions (e.g. some are discussed in Godfrey et al.). Could the authors discuss this (and if true include it in their discussion about limitations?).

Questions

- As far as I know, Proposition 3.4. and Proposition 4.4 follow from the characterisation of intertwiner groups done by Godfrey et al., NeurIPS’22 (e.g. see Theorem E.14 and Proposition 3.4), but the authors provide alternative proofs. In case these proofs are of independent technical interest, I would recommend that the authors briefly mention a proof sketch in the main paper and the differences with Godfrey et al. (this would also help in properly accrediting this prior work). - L210: The authors mention “It is natural to ask whether the group $G$ is still maximal in the other case”. What do they mean by the “other case”? - Perhaps an illustrative example of the resulting weight-sharing would help. Or maybe give some intuition on the resulting layers? - Maybe adding the exact number of parameters along with the performance in the tables would help in grasping the actual reduction. **Minor**. - L126-L128 are unclear. Perhaps describe in more detail the notation $\text{Aut}(\Delta_n)$ and the terms conjugation, semi-direct product etc. (depending on the importance of these statements). - L242: hyperparams --> Perhaps you mean learnable params?

Rating

7

Confidence

5

Soundness

4

Presentation

4

Contribution

3

Limitations

Some of the limitations of this work are mentioned in the text (last paragraph of Section 5 and conclusion section). However, I think that some have not been adequately discussed (see weaknesses). I would recommend adding a separate section for this purpose. I do not foresee any negative societal impact.

Summary

This paper studies the extension of permutation equivariant neural functionals to accommodate the monomial group, which is a generalization of the permutation group. This extension leads to a new class of NFN called monomial NFN that can also handle the scaling symmetry of positively homogenous activation (RELU) and sign-flipping symmetry of activations such as tanh. The paper then characterizes the subset of all monomial matrices that are either *preserved* by either ReLU or sin/tanh ($\sigma A(x) = A(\sigma x)$). The paper then proceeds to construct group $G_\mathcal{U}$ which is the product of all monomial groups that act on a weight space object $U$, and two subgroups of $G_\mathcal{U}$ under which the models with ReLU or sin / tanh are invariant. Finally, the paper constructs an affine weight space layer that is $G$-equivariant and also $G$-invariant weight space layers, which are more parameter efficient than prior works. Generalization prediction experiments show that the proposed layer performs better than prior works under standard conditions and significantly outperforms when the models are perturbed by scale, corroborating the statement made by the paper. On INR classification and editing tasks, monomial-NFNs are competitive or better than prior works.

Strengths

Overall, this is a well-written and technically solid paper with a novel contribution to the neural functional literature. The theoretical results seem sound. While I did not thoroughly check every proof detail, Figure 1 demonstrates that the constructed layer is indeed equivariant to the scale of the weights. The other experimental results are also quite compelling. The parameter saving could also potentially be a big upside of the monomial-NFN which can bring better generalization.

Weaknesses

I don't see major weaknesses in the paper. Perhaps one shortcoming is that the majority of performance improvement comes when the weights are scaled, which is not a very natural perturbation in reality. On the other hand, many tasks that cannot be "solved" by previous NFNs remain difficult for monomial NFNs (e.g., CIFAR10 classification), which might indicate some fundamental limitation of this line of work.

Questions

1. How do the computation cost and memory compare to NFNs? In my opinion, the issue with NFNs' big parameter count is that it's very memory-intensive. While monomial-NFN saves parameters, the benefit is perhaps less significant if it doesn't come with a computational advantage.

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

Authors have adequately discussed the limitation though it would be good to include some computation cost analysis.

Summary

- Paper improves neural functional networks by taking into account weight scaling properties of ReLU networks and weight sign flipping symmetries of sin or Tanh networks. - Monomial matrices are used to represent these symmetries, both permutation and scale/sign-flipping transformations. - Proposed model has fewer trainable (independent) parameters compared to the original NFN family of models.

Strengths

- Proposes a principled way to incorporate activation functions in neural network representations for MLPs and CNNs.

Weaknesses

- No discussion/comparison with previous works on the subject [1,2,3,5]. - No discussion of extensions to architectures with branches/ transformers? - Not the first to consider activation functions in representing neural network weights as claimed. In [1] and [2], activation functions are encoded as nodes in a graph. - The performance gain over the baselines is minimal. Perhaps this suggests that the role of activation function encoding for this task is minimal as shown in the original neural functional works and those of [1, 2, 3, 4]. - There is a focus on very specific activation functions (relu, sin, tanh), admittedly, these are common in many architectures, but the authors do not provide any discussion of how the proposed method can be applied to other activation functions.

Questions

- Minimal equivariance is the goal for this task as equivariance is generally easy to achieve. How does the proposed model satisfy minimal equivariance while ignoring permutations that are not functionally equivarant? - Can the model handle a mixed modelzoo with different activation functions? Note that [1] can handle this case and [3] already demonstrates that this is possible. Per my understanding, the proposed model will need different instantiations of the model for each modelzoo with a different activation function. - How are the input/output layers of the considered networks handled in the proposed framework? Note that these layers do not follow the permutation symmetries of MLPs(only the hidden layers do). ## References [1] Lim, Derek, et al. "Graph metanetworks for processing diverse neural architectures." arXiv preprint arXiv:2312.04501 (2023). [2] Kofinas, Miltiadis, et al. "Graph neural networks for learning equivariant representations of neural networks." arXiv preprint arXiv:2403.12143 (2024). [3] Andreis, Bruno, Soro Bedionita, and Sung Ju Hwang. "Set-based neural network encoding." arXiv preprint arXiv:2305.16625 (2023). [4] Unterthiner, Thomas, et al. "Predicting neural network accuracy from weights." arXiv preprint arXiv:2002.11448 (2020). [5] Zhou, Allan, Chelsea Finn, and James Harrison. "Universal neural functionals." arXiv preprint arXiv:2402.05232 (2024).

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

While limitations are not discussed, i do not see any immediate negative societal impact.