🤖 machine learning

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

This paper establishes a systematic theoretical framework proving the equivalence and universality of permutation-equivariant weight-space networks in both weight and function spaces, while demonstrating that slight model modifications based on these insights yield a 34% performance improvement over state-of-the-art methods.

Original authors: Adir Dayan, Yam Eitan, Haggai Maron

Published 2026-06-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Adir Dayan, Yam Eitan, Haggai Maron

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive library of recipes (neural networks). Usually, when we want to learn from these recipes, we just look at the final dish they produce (the output). But this paper is about a different kind of chef: one who studies the ingredients and the instructions themselves (the weights and parameters) to understand, predict, or modify the recipe.

This field is called Weight-Space Learning. The problem is that recipes have a weird quirk: you can swap the order of the steps (like mixing eggs before flour vs. flour before eggs) or rename the bowls, and the final dish tastes exactly the same. This is called symmetry.

To handle this, researchers built special "symmetry-aware" chefs (neural networks) that respect these rules. But the authors of this paper asked a big question: Are these chefs actually smart enough to do anything we ask them to do, or are they limited by their strict rules?

Here is the breakdown of their findings, using simple analogies:

1. The "Same Chef" Discovery

The paper looked at several different types of these "symmetry-aware" chefs (named DWS, GMN, NFN, etc.). They are built differently, like one uses a hammer and the other uses a screwdriver.

The Finding: The authors proved that despite looking different, all the top chefs are actually equally powerful. If one can solve a puzzle, they all can. It doesn't matter which specific "symmetry-aware" tool you pick; they all have the same "brainpower."
The Exception: There was one special chef (NFT) that seemed slightly different, but the paper found that if the recipes are "normal" (not weirdly degenerate), this chef is just as powerful as the rest.

2. The Four Types of Tasks

The authors categorized the jobs these chefs might be asked to do into four buckets. Think of these as different ways to interact with a recipe book:

Bucket A: The Taste Tester (Function-Space Functionals)
- Task: "Look at this recipe and tell me how tasty the final dish will be."
- Result: Perfect. These chefs can predict the outcome of any recipe perfectly, as long as the recipe isn't broken. They are universal for this.
Bucket B: The Ingredient Counter (Permutation-Invariant Functionals)
- Task: "Count the total amount of sugar in the recipe, regardless of which bowl it's in."
- Result: Mostly Perfect, with a catch. If the recipe has unique ingredients (e.g., every bowl has a different amount of sugar), the chefs are perfect. But if the recipe has "degenerate" cases (e.g., two bowls have exactly the same amount of sugar), the chefs might get confused and fail to distinguish them. This is a rare edge case, like a recipe where two steps are identical.
Bucket C: The Recipe Transformer (Function-Space Operators)
- Task: "Take this recipe for a small cake and turn it into a recipe for a giant cake."
- Result: The Limitation. Here is the big problem. If the input recipe is for a small cake, the chef is forced to output a recipe for a small cake (same architecture). They cannot magically create a "bigger" recipe if the original one wasn't built to handle it. It's like trying to fit a whole elephant into a shoebox; the box (the output architecture) is too small.
- The Fix: The authors realized that if you let the chef output a bigger recipe (a larger network) than the one they received, they can do the job perfectly.
Bucket D: The Recipe Editor (Permutation-Equivariant Operators)
- Task: "Take this recipe and tweak the steps, but keep the symmetry rules."
- Result: Similar to Bucket B. They are perfect as long as the recipes aren't in those weird "degenerate" states where ingredients are identical.

3. The Practical Solution: "Output Capacity Expansion"

Because of the limitation in Bucket C (the "Recipe Transformer" problem), the authors proposed a simple, clever trick called Output Capacity Expansion (OCE).

The Analogy: Imagine you ask a chef to bake a cake, but you only give them a small pan. They can't make a big cake. The authors say: "Don't give them one small pan. Give them eight small pans, bake eight tiny cakes, and then mix them together."
The Result: By having the model predict multiple networks and averaging them, they effectively create a "bigger" output capacity without changing the model's complexity.
The Proof: When they tested this on a standard benchmark (editing images represented as neural networks), this simple trick improved the results by 34% over the previous best methods.

Summary

The paper builds a theoretical map of what these "symmetry-aware" networks can and cannot do.

They are all equal: Different designs are just different flavors of the same powerful tool.
They are mostly perfect: They can handle almost any task, provided the input data isn't weirdly repetitive.
They hit a wall on "growing" tasks: They struggle to turn a small network into a big one unless you let them output a bigger structure.
The fix works: A simple trick to let them output multiple networks solves the problem and significantly boosts performance in the real world.

The authors conclude that by understanding these theoretical limits, we can design better tools for editing and analyzing neural networks, rather than just guessing.

Problem Statement
Weight-space learning investigates neural architectures that operate directly on the parameters of other neural networks, treating weights as structured data for tasks such as accuracy prediction, meta-optimization, and model editing. A key property of Multi-Layer Perceptron (MLP) weight spaces is their inherent symmetry: permuting neurons within a hidden layer alters the raw parameters but leaves the realized function unchanged. Consequently, state-of-the-art (SOTA) weight-space networks employ permutation-equivariant designs to respect this symmetry. However, such symmetry constraints may restrict the hypothesis space, potentially limiting the expressive power of these networks. While prior works have established specific capabilities (e.g., simulating forward passes), a comprehensive theoretical characterization of the expressive power of permutation-equivariant weight-space networks remains missing. Furthermore, unlike other geometric deep learning domains, weight-space learning involves maps operating on both weight space and function space, making expressivity analysis particularly subtle.

Methodology
The authors develop a systematic theoretical framework to analyze the expressivity of permutation-equivariant weight-space networks operating on MLP weights. The methodology proceeds through four main stages:

Architectural Equivalence: The paper first establishes that prominent permutation-equivariant architectures—including Deep Weight Space (DWS), Neural Functional Networks (NFN), Graph Meta-Networks (GMN), and Neural Graph GNNs (NG-GNN)—are equivalent in expressive power. This is proven by demonstrating that each architecture can approximate the base layers of the others, forming a cycle of expressive containment. Neural Functional Transformers (NFTs) are initially identified as an exception due to their non-standard attention mechanisms but are shown to match the expressive power of the others under a "general position" (GP) assumption.
Approximation Framework: The authors define four fundamental approximation settings to categorize the target maps:
- Function-space functionals: Maps from functions to vectors (e.g., accuracy prediction).
- Permutation-invariant functionals: Maps from weights to vectors that are invariant to neuron permutations but may depend on specific parameterizations (e.g., weight $\ell_2$ -norm).
- Function-space operators: Maps from functions to functions (e.g., model editing).
- Permutation-equivariant operators: Maps from weights to weights that respect symmetries but are not solely determined by the realized function (e.g., pruning masks).
Universality Analysis: For each setting, the paper analyzes conditions under which universal approximation is achievable. A critical component of this analysis is the "general position" (GP) assumption, which requires that bias terms within each hidden layer are pairwise distinct. This assumption excludes a measure-zero set of degenerate configurations.
Theoretical Proofs: The authors utilize tools from geometric deep learning, including separation properties, canonization maps, and the Stone-Weierstrass theorem, to prove universality results. They demonstrate that while universality fails in full generality for certain settings (specifically permutation-invariant functionals and operators on fixed architectures), it is recovered under the GP assumption or by allowing the output architecture to be sufficiently large.

Key Contributions

Expressive Equivalence: The paper proves that all prominent permutation-equivariant weight-space networks for MLPs (DWS, NFN, GMN, NG-GNN) possess identical expressive power. NFTs are shown to be equivalent to this class under the GP assumption.
Unified Approximation Framework: The authors identify and formalize four natural approximation settings, providing a structured lens through which to evaluate weight-space networks.
Universality Characterization:
- Function-space functionals: Proven to be universally approximable by invariant weight-space networks on any compact set.
- Permutation-invariant functionals: Shown not to be universally approximable in full generality (due to limitations in distinguishing non-equivalent weight configurations with identical computational graphs). However, universality is achieved under the GP assumption.
- Function-space operators: Proven to be non-universal when constrained to output networks of the same fixed architecture (due to bounded representational capacity). Universality is achieved if the output architecture is allowed to be sufficiently large to accommodate the target function's complexity.
- Permutation-equivariant operators: Similar to invariant functionals, universality fails in full generality but holds under the GP assumption.
Empirical Validation via Output Capacity Expansion (OCE): Guided by the theoretical finding that function-space operators require larger output capacities, the authors propose a simple modification: predicting an ensemble of $k$ output networks and averaging them. This increases the effective output capacity without adding parameters to the model itself.

Results

Theoretical: The paper provides a complete map of the expressivity landscape, clarifying exactly when and why weight-space networks succeed or fail in approximating specific types of maps. It identifies the "exclusion set" (degenerate bias configurations) as the primary barrier to universality in invariant settings.
Empirical: On the MNIST Implicit Neural Representation (INR) dilation benchmark (a standard function-space operator task), the proposed OCE modification yields significant improvements. Specifically, applying OCE to DWS and GMN architectures achieves a 34% improvement over the previous SOTA (ScaleGMN + GradMetaNet++). The results demonstrate that increasing the effective output capacity directly addresses the theoretical limitations identified in the function-space operator setting.

Significance
The paper claims to provide the first comprehensive theoretical foundation for the expressive power of permutation-equivariant weight-space networks. By unifying diverse architectural families into a single expressivity class and precisely characterizing the conditions for universality, the work offers principled guidance for designing weight-space models. The identification of output capacity as a bottleneck for function-space operators leads to a practical, theory-driven improvement (OCE) that achieves SOTA performance. The authors position this work as a necessary step to move beyond task-specific empirical results toward a deeper mathematical understanding of learning in weight spaces, distinguishing it from prior works that only established partial capabilities like forward pass simulation.

1. The "Same Chef" Discovery

2. The Four Types of Tasks

3. The Practical Solution: "Output Capacity Expansion"

Summary

More like this