Interfaces and Modules: A Scalable Representation for Weight-Space Learning across Architectures

Michał Kozyra; Michael Bronstein

cs.LGOct 15, 2026

Interfaces and Modules: A Scalable Representation for Weight-Space Learning across Architectures

Authors: Michał Kozyra, Michael Bronstein

Organizations: University of Oxford · AITHYRA

Abstract

Weight-space learning across architectures requires representations that respect parameter symmetries and scale to large networks. Existing equivariant tensor encoders require architecture-specific symmetry specifications, while neuron graphs make global attention costly. We introduce a typed graph organized around shared coordinate-permutation actions. Its nodes, called interfaces, represent coordinate spaces whose bound tensor axes must be permuted jointly; operation occurrences, called modules, form typed hyperedges. A compiler derives this graph from supported PyTorch models, with exact recovery of the permutation factors defined by its constraints. The Universal Weight Encoder (UWE) combines local equivariant updates to parameter features with global attention over invariant interface summaries. The small number of interfaces makes global attention practical, and a single encoder processes different supported architectures. UWE is competitive with architecture-specific state-of-the-art methods on established benchmarks. On a new zoo of 1,200 transformers spanning 100 architectures, it outperforms the evaluated baselines at ranking networks within unseen architectures, including architectures larger than those used for training, while using substantially less memory.

Figures & tables

Figure 1: The compiler derives interfaces and parameter bindings from network code. UWE processes the resulting graph, and a supervised loss trains the encoder and readout.

Panel of Figure 2 — Figure 2: Conventional diagrams (left) and interface–module graphs (right). Matching colours identify corresponding coordinate spaces; shades distinguish heads. Objects sharing a coordinate-permutation action map to one interface. Operations whose inputs and outputs share an interface become self-loops. The compiler (Section 4 ) constructs these graphs for the encoder (Section 3 ).

Network	Parameters	$A$	$S$	Interface widths
ResNet-18 (20 convolutions)	11.7M	12	2,880	64–512, three of each
ResNeXt-50 $32{\times}4$ d (53 convolutions)	25.0M	1,029	19,008	mostly 4–32
ViT-S/16	22.1M	157	28,032	384, 64, 1536
ViT-B/16	86.6M	301	56,064	768, 64, 3072
ViT-L/16	304.3M	793	148,480	1024, 64, 4096
GPT-2 Small	124.4M	301	56,064	768, 64, 3072

Table 1: Permutation-bearing interfaces

A

and coordinates

S

; fixed interfaces excluded. A transformer with a permutation-bearing residual interface, width

D

h

heads and

L

blocks has

A=1+L(2h+1)

and

S=D+L(2D+D_{\mathrm{ff}})

: one residual interface of width

D

, two interfaces of width

D/h

per head per block, and one feed-forward interface of width

D_{\mathrm{ff}}

per block. Configurations: Appendix H.4 .

Panel of Figure 3 — Figure 3: One layer of the encoder. Every module bound to interface $I$ has an axis of exactly $n_{I}$ cells, one per coordinate, whether rectangular ( $u$ , $u^{\prime}$ ) or diagonal ( $u^{\prime\prime}$ ); emission pools a module’s other axis into one message per coordinate, and delivery broadcasts the coordinate states back along it. The global pass runs over one summary per interface and its result is broadcast to all coordinates.

	INR accuracy (%)		CNN zoo
Method	MNIST	F-MNIST	$\tau$
DWSNet	$85.71{\scriptstyle\,\pm 0.57}$	$67.06{\scriptstyle\,\pm 0.29}$	–
NFN NP	$78.50{\scriptstyle\,\pm 0.23}$	$68.19{\scriptstyle\,\pm 0.28}$	$0.922{\scriptstyle\,\pm 0.001}$
NFN HNP	$79.11{\scriptstyle\,\pm 0.84}$	$68.94{\scriptstyle\,\pm 0.64}$	$0.934{\scriptstyle\,\pm 0.001}$
NG-GNN	$94.70{\scriptstyle\,\pm 0.30}$	$74.20{\scriptstyle\,\pm 0.40}$	$0.931{\scriptstyle\,\pm 0.002}$
NG-T	$97.30{\scriptstyle\,\pm 0.20}$	$74.80{\scriptstyle\,\pm 0.90}$	$0.935{\scriptstyle\,\pm 0.000}$

Table 2: Test accuracy for INR classification and Kendall’s

\tau

for CIFAR-10-GS CNN-zoo generalisation prediction. Higher is better; baseline results are as published.

		common, common			common, beyond		full, full
Model	Params	MSE	$\tau$	h	MSE	$\tau$	MSE	$\tau$	h
Predict the training mean	—	0.9720	—	—	1.5510	—	1.2860	—	—
NG-T, node 16, edge 4	72k	0.9519	0.2165	23.1 H	—	—	—	—	—
Transformer-NFN, $(c=4)$	1.18B	0.8206	0.3420	2.9 H	—	—	—	—	—
GMN, width 64	0.19M	0.9594	0.3766	1.5 H	1.4273	0.4015	—	—	—
NG-GNN, released readout	0.20M	0.5263	0.5022	2.7	1.0320	0.5795	0.4007	0.6990	13.7

Table 3: Test results by (training cohort, test cohort), at checkpoints selected by validation MSE. MSE averages errors over 60 training-standardised profile coordinates;

\tau

ranks networks by mean profile loss within each architecture, then averages over architectures. h: runtime in hours; H : H100, otherwise A10. Common, beyond reuses the common, common models, with no separate runtime reported. —: the model does not fit the cohort’s architectures. Results use one metamodel training seed per configuration. Variability across training seeds was not estimated due to computational constraints.

		common		beyond
Variant	Params	MSE	$\tau$	MSE	$\tau$
UWE	73k	0.28	0.71	0.54	0.72
Without codes	67k	0.27	0.67	0.75	0.66
Without global attention	56k	0.73	0.33	1.15	0.44
Without either	50k	0.84	0.38	1.40	0.05

Table 4: Common-trained ablations evaluated at their best-validation-MSE checkpoints. Metrics follow Table 3 .

		UWE			NG-GNN			GMN			Transformer-NFN
Target	$P$	0.2M	0.9M	1.8M	0.2M	0.9M	1.8M	0.2M	0.9M	1.8M	$c{=}2$	$c{=}10$
MLP	1M	0.46	0.54	0.59	2.28	5.11	7.21	3.81	8.19	11.39	–	–
MLP	30M	10.34	10.70	10.97	66.16	OOM i	OOM i	OOM t	OOM i	OOM i	–	–
CNN	10M	0.55	0.63	0.69	2.63	5.80	8.09	37.47	OOM t	OOM i	–	–
CNN	100M	4.30	4.54	4.72	25.47	56.77	OOM t	OOM i	OOM i	OOM i	–	–
Residual CNN	100M	4.20	4.52	4.79	25.49	56.82	OOM t	–	–	–	–	–

Table 5: Peak allocated memory (GiB) for a metamodel training step, including the first AdamW step: one H100, 75 GiB cap, batch one. Targets have random weights; predictive performance is not evaluated. NG-GNN and GMN match UWE’s trainable parameter count within 2.1%. Transformer-NFN is unmatched: at the 3M target,

c=2

and

c=10

have 381M and 1.924B parameters. OOM i /OOM t : observed inference/training OOM; –: unsupported or not applicable.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Permutation symmetry acts on coordinate spaces and carries the parameters along the axes bound to them. Colour marks the two exchanged coordinates.

Symmetry	Origin	Group	In $\Gamma$
hidden-unit permutation	coordinatewise activation	$S_{n}$	yes
channel permutation	convolution on channel axes	$S_{c}$	yes
within-group permutation	grouped-convolution layout	$\prod_{j}S_{c/g}$	yes
head exchange	exchange of complete heads	$S_{h}$	no; included in $G$
positive rescaling	positively homogeneous $\sigma$	$\mathbb{R}_{>0}^{n}$	no
sign flip	odd $\sigma$	$\{\pm 1\}^{n}$	no

Appendix

Table 6: Parameter symmetries and their treatment in this paper. The last column refers to the within-interface group

\Gamma

. A permutation action is included only when the corresponding interfaces are permutation-bearing and the model’s bindings permit it. Permutation subgroup means that

\Gamma

contains the permitted coordinate permutations within the stated continuous group, without asserting the remaining transformations. Supported head exchanges belong to the enlarged group

G=\Gamma\rtimes E

. The attention product expressions describe bias-free projections; affine projections require the corresponding bias transformations.

Architectures	100 distinct $(d,L,h,r)$
Parameters	112,676 to 9,061,732; median 1.05M
Width $d$	64 to 608
Depth $L$	2 (30), 3 (31), 4 (23), 6 (16)
Heads $h$	2 to 16; head width 32 to 96
FFN ratio $r$	2 (34), 3 (38), 4 (28)
Runs	1,200: one per recipe per architecture

Appendix

Table 7: Transformer zoo summary. Parentheses give architecture counts at each depth or feed-forward ratio.

Figure 5: Transformer zoo labels. (a) Mean profile loss versus parameter count, coloured by training-data fraction; the fitted slope is

-0.14

nats per decade. (b) Distributions of per-condition mean loss. (c) Range of mean profile loss across the 12 runs of each architecture; the dashed line marks the median,

0.83

nats.

Setting	MNIST-INR	Fashion-MNIST-INR	CNN zoo
$d_{E}$	16	16	40
$d_{H}$	80	80	160
Layers	4	4	8
Attention heads	4	4	8
Global tokens	4	4	4
Learned probes	64	64	0

Appendix

Table 8: Standard-benchmark settings. Counts include the encoder, readout and probes. Hyperparameters were loosely chosen to align the model with baseline methods.

Experiment	$d_{E}$	$d_{H}$	$g$	Layers	Readout	Parameters
Loss profile, small	16	32	0	4	$H$	72,636
Loss profile, large	32	256	0	4	$H$	3,458,684
Budget ladder, small	16	56	4	4	$H,G$	190,057
Budget ladder, medium	16	128	4	4	$H,G$	906,529
Budget ladder, large	16	180	4	4	$H,G$	1,759,277
Native ViT/GPT-2	4	128	4	4	$H,G$	887,881

Appendix

Table 9: UWE configurations for loss-profile prediction and resource measurements;

g

is the number of global tokens. Readout uses pooled interface features

H

and, when present, global-token features

G

. Counts include the encoder and prediction head. The parameter-matched native comparison uses

d_{E}=4

Metamodel	Trainable parameters
UWE, $d_{E}=16$ , $d_{H}=32$	72,636
UWE, $d_{E}=32$ , $d_{H}=256$	3,458,684
NG-GNN, output-node concatenation	203,516
NG-GNN, all-node mean	102,140
NG-T, node 16, edge 4	72,388
GMN, width 64	185,340

Appendix

Table 10: Headline parameter counts, including input encodings and prediction heads. Small UWE approximately matches NG-T; the full comparison is not parameter-matched.

Method	Params	H100	A10	0.1–0.3M	0.3–1M	1–3M	3–10M	(s)
		Fits, of 85		Peak GiB by target size, H100				Step
UWE, $d_{E}=32$ , $d_{H}=256$	3.46M	85	85	0.72	1.43	3.24	9.17	0.81
UWE, $d_{E}=16$ , $d_{H}=32$	73k	85	85	0.22	0.46	1.13	3.19	0.82
NG-GNN, width 32	0.20M	85	85	0.54	1.60	4.32	12.91	0.10
GMN, width 64	0.19M	85	69	1.66	5.13	14.58	44.14	0.15
NG-T, node 16, edge 4	72k	52	26	21.93 17	57.54	68.47 12	–	1.47

Appendix

Table 11: Zoo resource use: Configurations follow Table 10 with batch one. NG-GNN uses the released output-node concatenation readout. The 85 architectures span widths 64–608 and 0.11M–8.3M parameters. Fits counts completed steps under 70 GiB H100 and 20 GiB A10 caps. Peak is the maximum over successful H100 cases per bin (18/23/25/19 architectures); superscripts give partial-coverage counts. Step is the median over each method’s successful H100 cases. GMN and NG-T failures are OOMs except one NG-T implementation failure; Transformer-NFN coverage is width-limited. –: no completed measurement; NR: not evaluated.

Target	$P$ (M)	Encoder	Infer. (s)	Infer. (GiB)	Step (s)	Step (GiB)
ViT-B/16	86.57	UWE, $d_{E}=4$	2.66	3.64	17.95	17.72
		UWE, $d_{E}=8$	2.54	5.93	17.37	24.35
		UWE, $d_{E}=16$	2.57	11.12	17.66	37.39
		NG-GNN	OOM	–	–	–
		GMN	OOM	–	–	–
ViT-L/16	304.33	UWE, $d_{E}=4$	7.83	12.49	54.60	58.19

Appendix

Table 12: Native-model resource measurements: one H100, 75 GiB cap, FP32, batch one, four rounds and checkpointing. UWE has 0.888M/0.894M/0.907M parameters at

d_{E}=4/8/16

, with

d_{H}=128

. Graph baselines match

d_{E}=4

within 2%. Memory includes the initial forward or AdamW step; time is the median of five repetitions. OOM: failure in the indicated phase; –: no completed measurement. Random source weights are used; predictive performance is not measured.

Figure 6: Interface inference for concatenation and addition. Concatenation splits $c$ at position 2; addition propagates this split to $s$ and $h_{3}$ . The resulting atoms form two interfaces with permutation group $S_{2}\times S_{3}$ . Colors indicate their bindings to rows of $W_{3}$ and columns of $W_{o}$ . Input and output coordinates remain fixed.

Operation	Constraints
Linear layer without parameter sharing	No input–output glue.
Elementwise operation; addition or multiplication with aligned coordinate axes	Glue the participating coordinate axes.
Concatenation along the coordinate axis	Glue each input to its output interval and split at the interval boundaries.
Concatenation along another axis	Glue each input coordinate axis to the output coordinate axis; no coordinate splits.
Depthwise convolution with $g=n_{\rm in}=n_{\rm out}>1$	Glue input and output coordinate axes.
Other grouped convolution, $g>1$	Split input and output coordinate axes at group boundaries; no input–output glue.

Appendix

Table 13: Representative operator rules. Each row assumes a supported layout. Fixed-axis and parameter-sharing constraints may further restrict the permitted permutations.

Representation	Permutation action	Order	Non-symmetries
Reference subgroup	$G_{\mathrm{att}}$	663,552	0
QKV-triple graph ( Kofinas et al., 2024 )	$S_{8}$	40,320	39,168
With fixed head labels ( Lim et al., 2024 )	$S_{4}\times S_{4}$	576	0
Transformer-NFN ( Tran et al., 2025 )	$G_{\mathrm{att}}$	663,552	0
UWE	$G_{\mathrm{att}}$	663,552	0

Appendix

Table 14: Internal permutation actions for two heads of width four, with stream coordinates fixed. Non-symmetries are listed transformations that do not preserve the attention function for all parameter values. Transformer-NFN also supports continuous basis changes, omitted from these counts.

Operation	Permutation constraints	State bindings
Linear layer; convolution with one group	No input–output glue.	Rectangular weights; biases bind the output. Kernel offsets are folded.
Depthwise convolution, $g=n_{\mathrm{in}}=n_{\mathrm{out}}>1$	Glue input and output.	Diagonal kernels and biases; kernel offsets are folded.
Other grouped convolution, $g>1$	Split input and output at group boundaries; no input–output glue.	Rectangular kernels connect matching groups; biases bind the output.
BatchNorm; single-axis LayerNorm	Glue input and output.	Diagonal affine parameters. BatchNorm running statistics follow the channel action.
Pointwise activation, spatial pooling, evaluation-mode dropout	Glue input and output.	No stored state.
Aligned addition or multiplication	Glue both operands and output.	No stored state.

Appendix

Table 15: Implemented constraint patterns and state bindings. Each row is subject to the restrictions below. Fixed-axis and sharing constraints may further restrict the action. Biases and shared state are stored without duplication.

Weight-Space Learning

See all 17 papers

Oct 1, 2026

CrossGMN: Graph Metanetworks for Cross-Architecture Weight-Space Transformations

Weight-space networks operate directly on parameters of other neural networks, enabling tasks such as predicting model properties, editing trained models, and generating weights. Weight-space symmetries such as neuron permutations make equivariance a key design principle. However, existing equivariant weight-space architectures have primarily been studied for transformations that preserve the network architecture. In contrast, many practical transformations, including model compression and upscaling, map a trained source network into a target network with a different architecture. In this setting, the source and target permutation symmetries act on different parameter spaces, making equivariance less straightforward to formulate. Our key idea for addressing this mismatch is to reformulate cross-architecture operators with two inputs: a trained source network and an initialization of the target network. This lets us define equivariant cross-architecture operators that refine the initialization of the target network using information from the source network, while being invariant to source-network permutations and equivariant to target-network permutations. Based on this formulation, we introduce CrossGMN, a graph metanetwork that jointly processes both networks through symmetry-preserving cross-network message passing. We prove CrossGMN is universal for continuous cross-architecture operators on compact sets under a general-position assumption. We evaluate CrossGMN for model compression, predicting a smaller network's parameters to accelerate subsequent knowledge distillation. Across 2-D and 3-D INRs and image classification with MLPs, CNNs, and Vision Transformers, CrossGMN speeds up distillation by up to 8.89x, transfers across datasets without retraining (3.78x), and a single model can accelerate compression from heterogeneous source architectures into a common target architecture.

Adir Dayan, Yam Eitan, Haggai Maron

Technion – Israel Institute of Technology · NVIDIA Research

Model Compression Hypernetworks Weight-Space Analysis Equivariant GNNs

Ask PDF

Oct 13, 2026

Graph Metanetworks for Transformer Based Language Models

Weight-space learning treats the weights of neural networks as data, processed by other models called metanetworks. A key challenge is that different weights can realize the same function, related by a group of symmetries the metanetwork should respect. While symmetry-aware metanetworks for simple architectures such as MLPs are well developed, a comprehensive treatment of transformers, and in particular transformer-based language models, is still missing. Importantly, their weight space has a richer symmetry group: beyond neuron permutations, attention introduces continuous general-linear symmetries of the attention matrices and permutations of entire heads. Existing methods for transformer weights either capture only part of these symmetries or are restricted to a single, fixed architecture, and none accounts for language-model-specific components such as the token embeddings. Building on graph metanetworks, we propose a graph-based approach that is both symmetry-aware and applicable across depths and widths. We formalize a weight-space symmetry group of complete decoder-only language models and provide a graph construction that is invariant to it. We further introduce two new transformer-population datasets with novel tasks, such as training-corpus prediction and fine-tuning detection, on which our approach notably outperforms existing baselines.

Yonathan Wolanowsky, Guy Bar-Shalom, Fabrizio Frasca +1

Technion – Israel Institute of Technology · NVIDIA Research

Weight-Space Learning Transformer Hypernetworks Decoder-Only Language Models

Ask PDF

Mar 10, 2026

A Survey of Weight Space Learning: Understanding, Representation, and Generation

Neural network weights are typically viewed as the end product of training, while most deep learning research focuses on data, features, and architectures. However, recent advances show that the set of all possible weight values (weight space) itself contains rich structure: pretrained models form organized distributions, exhibit symmetries, and can be embedded, compared, or even generated. Understanding such structures has tremendous impact on how neural networks are analyzed and compared, and on how knowledge is transferred across models, beyond individual training instances. This emerging research direction, which we refer to as Weight Space Learning (WSL), treats neural weights as a meaningful domain for analysis and modeling. This survey provides the first unified taxonomy of WSL. We categorize existing methods into three core dimensions: Weight Space Understanding (WSU), which studies the geometry and symmetries of weights; Weight Space Representation (WSR), which learns embeddings over model weights; and Weight Space Generation (WSG), which synthesizes new weights through hypernetworks or generative models. We further show how these developments enable practical applications, including model retrieval, continual and federated learning, neural architecture search, and data-free reconstruction. By consolidating fragmented progress under a coherent framework, this survey highlights weight space as a learnable, structured domain with growing impact across model analysis, transferring, and weight generation. We release an accompanying resource at https://github.com/Zehong-Wang/Awesome-Weight-Space-Learning.

Xiaolong Han, Zehong Wang, Bo Zhao +8

University of Surrey · University of Notre Dame · University of California San Diego +4

Weight-Space Learning Weight-Space Analysis

Ask PDF

Abstract

Weight-Space Learning

Explore similar work

CrossGMN: Graph Metanetworks for Cross-Architecture Weight-Space Transformations

Graph Metanetworks for Transformer Based Language Models

A Survey of Weight Space Learning: Understanding, Representation, and Generation