On the Expressive Power of Permutation-Equivariant Weight-Space Networks

TL;DR

This work addresses the expressivity of permutation-equivariant weight-space networks—architectures that operate directly on the parameters of other neural nets. It establishes a unified theory showing expressive-equivalence across major architectures and characterizes universality across four natural settings: function-space functionals, permutation-invariant functionals, function-space operators, and permutation-equivariant operators. Universality holds under mild general-position assumptions for several settings, with explicit edge cases where it fails for fixed-architecture inputs but can be recovered by increasing input capacity, as well as for certain operator classes. The results provide principled guidance for designing weight-space models, with implications for model editing, meta-learning, and weight-space transfer in pretrained-model regimes.

Abstract

Weight-space learning studies neural architectures that operate directly on the parameters of other neural networks. Motivated by the growing availability of pretrained models, recent work has demonstrated the effectiveness of weight-space networks across a wide range of tasks. SOTA weight-space networks rely on permutation-equivariant designs to improve generalization. However, this may negatively affect expressive power, warranting theoretical investigation. Importantly, unlike other structured domains, weight-space learning targets maps operating on both weight and function spaces, making expressivity analysis particularly subtle. While a few prior works provide partial expressivity results, a comprehensive characterization is still missing. In this work, we address this gap by developing a systematic theory for expressivity of weight-space networks. We first prove that all prominent permutation-equivariant networks are equivalent in expressive power. We then establish universality in both weight- and function-space settings under mild, natural assumptions on the input weights, and characterize the edge-case regimes where universality no longer holds. Together, these results provide a strong and unified foundation for the expressivity of weight-space networks.

Figures (4)

• Figure 1: Expressivity landscape for permutation-equivariant weight-space networks on MLPs (blue arrows: new contributions; red arrows: previous results). $\mathcal{V}$ and ${\mathcal{F}}$ denote weight-space and function-space, respectively. Left: Equivalence of permutation-equivariant networks. Center: Assuming general position, all weight-space networks are universal across all approximation settings, strengthening prior feed-forward simulation results for DWS and GMN. Right: Approximation settings in weight space.
• Figure 2: Real-world examples of target functions for all approximation settings in weight-space learning. $\mathcal{V}$ and ${\mathcal{F}}$ denote weight-space and function-space, respectively. 1. Lim2024GMNZhou2023NFNKofinas2024NeuralGraphs 2. Navon2023Zhou2023NFNLim2022SignBasis 4. gelberg2025gradmetanet 5. zhou2023neuralLim2022SignBasisKofinas2024NeuralGraphs 6. Navon2023 7. Zhou2023NFN 8. Kofinas2024NeuralGraphsgelberg2025gradmetanetzhou2024universal
• Figure 3: Computational graphs induced by weights ${\bm{v}}$ and ${\bm{v}}'$ used in the proof of Proposition \ref{['prop:invariant_nonuniv']}. The weight matrices have binary entries, where $1$ corresponds to an edge and $0$ to the absence of an edge, and bias terms are equipped with features encoding their layer index. Both graphs admit identical 1-WL colorings (node colors) and are therefore indistinguishable by message-passing GNNs (NG-GNNs). Formal definitions of ${\bm{v}}$ and ${\bm{v}}'$ are given in Appendix \ref{['sec:app:invariant_equivalence_hierarchy']}, Proposition \ref{['prop:app:DWS_not_G']}.
• Figure 4: Diagram illustrating Theorem \ref{['thm:func_to_func_univ_off_exclusion']}. A function-space operator $\Psi : K\subset {\mathcal{C}}(X,{\mathbb{R}}^m) \to {\mathcal{C}}(X,{\mathbb{R}}^m)$ is approximated by a permutation-equivariant weight-space operator $\Phi : K'\subset \mathcal{V}_{A} \setminus {\mathcal{E}}_A \to \mathcal{V}_{A}$. The set $K'$ serves as an approximation of $K$ under the realization map ${\mathcal{R}}$. The diagram is approximately commutative, showing that $\Phi$ approximates $\Psi$ via ${\mathcal{R}}$.

Theorems & Definitions (110)

• Definition 3.1: weight space
• Definition 3.2: weight-space representation
• Definition 3.3: invariance and equivariance
• Definition 3.4: exclusion set
• Definition 4.1: approximation via weight-space maps
• Definition 5.1
• Theorem 5.2
• Proposition 5.3
• Theorem 6.1
• proof : Proof sketch.