Permutation-Equivariant Flow Matching for Alignment-Free Neural Weight Generation

2026-09-26

AI Summary

This study addresses the challenge of learning weight distributions in neural network generation, where neuron permutation symmetry complicates distribution modeling and conventional alignment methods incur high computational costs while depending on base models. To overcome these limitations, this work proposes a permutation-equivariant flow matching framework that parameterizes the velocity field via graph meta-networks. By directly learning weight distributions from ensembles of independently trained networks, it enables end-to-end neural weight generation without requiring explicit alignment. This approach transcends approximate alignment constraints, supporting generative generalization across heterogeneous architectures and unseen widths. Experiments demonstrate that the method effectively reproduces joint statistical properties and achieves performance comparable to ensemble learning under domain shift, thereby validating its efficacy in generating neural weights beyond mere memorization.

Abstract

A trained neural network can be represented by a parameter vector in high dimensions. Learning distributions over these vectors enables the generation of new models across various tasks and architectures. A central challenge is permutation symmetry: permuting hidden neurons can produce distant parameter vectors representing the same function. This introduces variations that a generative model must account for when learning from trained networks. Existing methods typically address this using networks derived from a common base model or costly approximate neuron alignment. We instead parameterize a flow-matching velocity field with a permutation-equivariant Graph Meta Network, enabling direct learning from independently trained networks without alignment. Extensive experiments show that our method closely reproduces the joint statistics of accuracy, functional similarity, and weight similarity of independently trained collections, providing evidence of generation beyond checkpoint memorization. A single conditional model also generates task-specific networks on heterogeneous architectures and generalizes to unseen hidden-width configurations. On a tabular domain-shift task, intermediate conditioning produces individual networks with performance comparable to logit ensembles across both domains. Taken together, our results show how permutation equivariance enables learning from diverse collections of independently trained networks without permutation alignment.

Problem

• neural weight generation
• permutation symmetry
• parameter distribution learning
• alignment-free

Innovation

• Permutation Equivariance
• Flow Matching
• Neural Weight Generation
• Graph Meta Network
• Alignment-Free Learning