Graph Metanetworks for Transformer Based Language Models

Yonathan Wolanowsky; Guy Bar-Shalom; Fabrizio Frasca; Haggai Maron

cs.LGOct 13, 2026

Graph Metanetworks for Transformer Based Language Models

Authors: Yonathan Wolanowsky, Guy Bar-Shalom, Fabrizio Frasca, Haggai Maron

Organizations: Technion – Israel Institute of Technology · NVIDIA Research

Abstract

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.

Figures & tables

Network Block	Weight Matrix	Group Action	Group Element
(1) Embed.	$\mathbf{E}$	$P_{V}^{\top}\,\mathbf{E}\,P_{d}$	$S_{V}\times S_{d}$
(2) Pre-attn.	$\bm{\gamma}^{a,\ell}$	$\bm{\gamma}^{a}P_{d}$	$S_{d}$
(2) norm	$\bm{\beta}^{a,\ell}$	$\bm{\beta}^{a}P_{d}$
(3) Attention	$\mathbf{W}^{q,h,\ell}$	$P_{d}^{\top}\,\mathbf{W}^{q,\sigma(h)}\,G_{qk}^{\sigma(h)}$	$S_{\wr}$
	$\mathbf{W}^{k,h,\ell}$	$P_{d}^{\top}\,\mathbf{W}^{k,\sigma(h)}\,(G_{qk}^{\sigma(h)})^{-\top}$
	$\mathbf{W}^{v,h,\ell}$	$P_{d}^{\top}\,\mathbf{W}^{v,\sigma(h)}\,G_{vo}^{\sigma(h)}$

Table 1: The symmetry group action over a weight space

w\in\mathcal{W}

. The attention group is

S_{\wr}\coloneqq S_{d}\times\bigl[\bigl(\mathrm{GL}(d_{k})\times\mathrm{GL}(d_{v})\bigr)\wr S_{H}\bigr]

. The action column omits the layer index

\ell

, which is left unchanged.

Panel of Figure 1 — Figure 1: The multi-head attention subgraph under the three graph constructions, for an attention layer with $d{=}4$ and $H{=}2$ heads ( $d_{k}{=}d_{v}{=}2$ ); the rest of GMT-GL and GMT-P is identical. Left: GMT-GL places the products $W^{(Q,i)}{W^{(K,i)}}^{\top}$ and $W^{(V,i)}W^{(O,i)}$ on the edges, and the head nodes $u^{h,1}$ bind the nodes of each head together. Middle: GMT-P gives each projection matrix its own edge set. Right: [ 20 ] and [ 18 ] which share the same construction for MHA, noted as GMN.

Flat weight vector baselines
Method	KL	MAE
DeepSets	$0.2681\pm 0.0002$	$0.1974\pm 0.0002$
MLP	$0.2637\pm 0.0067$	$0.1951\pm 0.0031$
MLP+Aug	$0.2715\pm 0.0009$	$0.1988\pm 0.0003$
Graph constructions
GMN	$0.1695\pm 0.0123$	$0.1439\pm 0.0064$

Table 2: Performance of all methods on the corpus-mixture prediction task ( Section 4.1 ).

Method	KL	MAE
NFN-T	$0.2942\pm 0.0060$	$0.2101\pm 0.0008$
GMT-P (ours)	$0.2040\pm 0.0121$	$\mathbf{0.1644}\pm 0.0080$
GMT-GL (ours)	$\mathbf{0.2038}\pm 0.0253$	$0.1652\pm 0.0126$

Table 3: Performance comparison against NFN-T [ 34 ] over a subset of

500

homogeneous architectures from the corpus-mixture dataset ( Section 4.1 ).

Figure 2: Test KL divergence on corpus-mixture prediction, evaluated as a function of training set size and across ablations of the (un)embedding components.

Method	ROC-AUC	Macro-F1	Accuracy
MLP	$0.6237$	$0.3333$	$0.5$
MLP+Aug	$0.5891$	$0.3333$	$0.5$
DeepSets	$0.5777$	$0.3333$	$0.5$
GMN	$0.9359$	$0.8370$	$0.8374$
GMT-P (ours)	$0.9495$	$0.8473$	$0.8475$
GMT-GL (ours)	$\mathbf{0.9557}$	$\mathbf{0.8611}$	$\mathbf{0.8613}$

Table 4: Results on fine-tuning detection ( Section 4.2 ).

Method	Macro-F1	ROC-AUC	Accuracy
MLP	0.4314	0.5606	0.7587
MLP+Aug	0.5260	0.7797	0.7593
DeepSets	0.4313	0.5502	0.7586
GMN	$0.9043$	$0.9743$	$0.9043$
GMT-P (ours)	$\mathbf{0.9520}$	$\mathbf{0.9926}$	$\mathbf{0.9520}$

Table 5: Weight-lv fine-tuning detection ( Section 4.3 ).

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Node layers	Weight edges	Parameter-free edges
(1) Embedding	$\mathcal{V}_{V}$ , $\mathcal{R}^{0}$	$\mathcal{E}_{\mathrel{\vbox{\hbox{\hbox to4.52pt{\vbox to3.79pt{\pgfpicture\makeatletter\hbox{\hskip 0.15pt\lower-0.63551pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.3pt} \lx@inpgf@ignorespaces{ {}{{}}{} {\lx@inpgf@ignorespaces}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 0 L 5.48 3.27}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{0.5126}{-0.5126}{0.85863}{4.09221pt}{2.44312pt}\lxSVG@begingroup@{transform=matrix(0.85863 0.5126 -0.5126 0.85863 5.66 3.38)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces {{\lx@inpgf@ignorespaces}{}}{{}}{} {{\lx@inpgf@ignorespaces}{}}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 3.49 L 5.48 0.21}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{-0.5126}{0.5126}{0.85863}{4.09221pt}{0.07689pt}\lxSVG@begingroup@{transform=matrix(0.85863 -0.5126 0.5126 0.85863 5.66 0.11)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces } \lxSVG@closescope {\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}}}}(\mathcal{V}_{V},\mathcal{R}^{0};\mathbf{E})$	–
(2) Pre-attn. norm	$\mathcal{V}_{\gamma,\beta}^{a,\ell}$	$\mathcal{E}_{\mathrel{\vbox{\hbox{\hbox to4.52pt{\vbox to3.79pt{\pgfpicture\makeatletter\hbox{\hskip 0.15pt\lower-0.63551pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.3pt} \lx@inpgf@ignorespaces{ {}{{}}{} {\lx@inpgf@ignorespaces}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 0 L 5.48 3.27}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{0.5126}{-0.5126}{0.85863}{4.09221pt}{2.44312pt}\lxSVG@begingroup@{transform=matrix(0.85863 0.5126 -0.5126 0.85863 5.66 3.38)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces {{\lx@inpgf@ignorespaces}{}}{{}}{} {{\lx@inpgf@ignorespaces}{}}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 3.49 L 5.48 0.21}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{-0.5126}{0.5126}{0.85863}{4.09221pt}{0.07689pt}\lxSVG@begingroup@{transform=matrix(0.85863 -0.5126 0.5126 0.85863 5.66 0.11)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces } \lxSVG@closescope {\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}}}}(\mathcal{V}_{\gamma,\beta}^{a,\ell},\mathcal{R}^{2\ell-2};[\bm{\gamma}^{a,\ell};\bm{\beta}^{a,\ell}])$	–
(3) Attention	$\mathcal{V}_{d}^{h,\ell}$ , $u^{h,\ell}$ , $\mathcal{R}^{2\ell-1}$	$\mathcal{E}_{\mathrel{\vbox{\hbox{\hbox to4.52pt{\vbox to3.79pt{\pgfpicture\makeatletter\hbox{\hskip 0.15pt\lower-0.63551pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.3pt} \lx@inpgf@ignorespaces{ {}{{}}{} {\lx@inpgf@ignorespaces}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 0 L 5.48 3.27}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{0.5126}{-0.5126}{0.85863}{4.09221pt}{2.44312pt}\lxSVG@begingroup@{transform=matrix(0.85863 0.5126 -0.5126 0.85863 5.66 3.38)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces {{\lx@inpgf@ignorespaces}{}}{{}}{} {{\lx@inpgf@ignorespaces}{}}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 3.49 L 5.48 0.21}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{-0.5126}{0.5126}{0.85863}{4.09221pt}{0.07689pt}\lxSVG@begingroup@{transform=matrix(0.85863 -0.5126 0.5126 0.85863 5.66 0.11)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces } \lxSVG@closescope {\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}}}}(\mathcal{R}^{2\ell-2},\mathcal{V}_{d}^{h,\ell};\mathbf{M}^{h,\ell})$	$\mathcal{E}_{\mathrel{\vbox{\hbox{\scalebox{0.75}{$ \rightrightarrows $}}}}}(\mathcal{R}^{2\ell-2},\mathcal{V}_{d}^{h,\ell})$
		$\mathcal{E}_{\mathrel{\vbox{\hbox{\hbox to4.52pt{\vbox to3.79pt{\pgfpicture\makeatletter\hbox{\hskip 0.15pt\lower-0.63551pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.3pt} \lx@inpgf@ignorespaces{ {}{{}}{} {\lx@inpgf@ignorespaces}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 0 L 5.48 3.27}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{0.5126}{-0.5126}{0.85863}{4.09221pt}{2.44312pt}\lxSVG@begingroup@{transform=matrix(0.85863 0.5126 -0.5126 0.85863 5.66 3.38)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces {{\lx@inpgf@ignorespaces}{}}{{}}{} {{\lx@inpgf@ignorespaces}{}}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 3.49 L 5.48 0.21}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{-0.5126}{0.5126}{0.85863}{4.09221pt}{0.07689pt}\lxSVG@begingroup@{transform=matrix(0.85863 -0.5126 0.5126 0.85863 5.66 0.11)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces } \lxSVG@closescope {\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}}}}(\mathcal{V}_{d}^{h,\ell},\mathcal{R}^{2\ell-1};\mathbf{N}^{h,\ell})$	$\mathcal{E}_{\mathrel{\vbox{\hbox{\scalebox{0.75}{$ \rightrightarrows $}}}}}(\mathcal{R}^{2\ell-2},\mathcal{R}^{2\ell-1})$
			$\{(u^{h,\ell},x^{h,\ell}_{i}):i\in[d]\}$
(4) Pre-MLP norm	$\mathcal{V}_{\gamma,\beta}^{b,\ell}$	$\mathcal{E}_{\mathrel{\vbox{\hbox{\hbox to4.52pt{\vbox to3.79pt{\pgfpicture\makeatletter\hbox{\hskip 0.15pt\lower-0.63551pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.3pt} \lx@inpgf@ignorespaces{ {}{{}}{} {\lx@inpgf@ignorespaces}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 0 L 5.48 3.27}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{0.5126}{-0.5126}{0.85863}{4.09221pt}{2.44312pt}\lxSVG@begingroup@{transform=matrix(0.85863 0.5126 -0.5126 0.85863 5.66 3.38)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces {{\lx@inpgf@ignorespaces}{}}{{}}{} {{\lx@inpgf@ignorespaces}{}}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 3.49 L 5.48 0.21}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{-0.5126}{0.5126}{0.85863}{4.09221pt}{0.07689pt}\lxSVG@begingroup@{transform=matrix(0.85863 -0.5126 0.5126 0.85863 5.66 0.11)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces } \lxSVG@closescope {\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}}}}(\mathcal{V}_{\gamma,\beta}^{b,\ell},\mathcal{R}^{2\ell-1};[\bm{\gamma}^{b,\ell};\bm{\beta}^{b,\ell}])$	–

Appendix

Table 6: Node layers and edge sets of

\Gamma_{\mathrm{GL}}(w)

, for every block

\ell\in[L]

and head

h\in[H]

, grouped by the network component they represent and numbered as in Table 1 . The residual layers

\mathcal{R}^{2\ell-2},\mathcal{R}^{2\ell-1},\mathcal{R}^{2\ell}

are the layers

\mathcal{V}_{d},\mathcal{V}_{d}^{\prime},\mathcal{V}_{d}^{\prime\prime}

of block

\ell

in Section 3.2 , which we denote here differently to conveniently account for the nodes of different blocks. The nodes are named

\mathcal{V}_{V}=\{v_{i}\}

\mathcal{V}_{V}^{\prime}=\{v^{\prime}_{i}\}

\mathcal{R}^{k}=\{r^{k}_{i}\}

\mathcal{V}_{d}^{h,\ell}=\{x^{h,\ell}_{i}\}

\mathcal{V}_{ff}^{\ell}=\{z^{\ell}_{i}\}

and

\mathcal{V}_{\gamma,\beta}=\{g,s\}

, with

u^{h,\ell}

the head node of head

h

. The attention products are

\mathbf{M}^{h,\ell}=\mathbf{W}^{q,h,\ell}(\mathbf{W}^{k,h,\ell})^{\top}

and

\mathbf{N}^{h,\ell}=\mathbf{W}^{v,h,\ell}\mathbf{W}^{O,\ell}_{h}

, and

[\bm{\gamma};\bm{\beta}]

is the

2\times d

matrix with rows

\bm{\gamma}

and

\bm{\beta}

. Fully-connected edge sets

\mathcal{E}_{\mathrel{\vbox{\hbox{\hbox to4.52pt{\vbox to3.79pt{\pgfpicture\makeatletter\hbox{\hskip 0.15pt\lower-0.63551pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.3pt} \lx@inpgf@ignorespaces{ {}{{}}{} {\lx@inpgf@ignorespaces}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 0 L 5.48 3.27}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{0.5126}{-0.5126}{0.85863}{4.09221pt}{2.44312pt}\lxSVG@begingroup@{transform=matrix(0.85863 0.5126 -0.5126 0.85863 5.66 3.38)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces {{\lx@inpgf@ignorespaces}{}}{{}}{} {{\lx@inpgf@ignorespaces}{}}{}{}{}{}{}{{}}\lxSVG@stroke\lxSVG@drawpath@unclipped{M 0 3.49 L 5.48 0.21}{fill:none} {{}{{}}{}{}{{}}{{{\lx@inpgf@ignorespaces}}{{{\lx@inpgf@ignorespaces}}{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{0.85863}{-0.5126}{0.5126}{0.85863}{4.09221pt}{0.07689pt}\lxSVG@begingroup@{transform=matrix(0.85863 -0.5126 0.5126 0.85863 5.66 0.11)} \lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke-dasharray=none,stroke-dashoffset=0.0pt} \lxSVG@begingroup@{stroke-linecap=round} \lxSVG@begingroup@{stroke-linejoin=round} \lxSVG@drawpath@unclipped{M -1.3 1.6 C -1.07 0.64 -0.54 0.19 0 0 C -0.54 -0.19 -1.07 -0.64 -1.3 -1.6}{fill:none} \lxSVG@closescope \lxSVG@closescope }}{{\lx@inpgf@ignorespaces}}}}\lx@inpgf@ignorespaces } \lxSVG@closescope {\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}}}}(\cdot,\cdot;\cdot)

consist of weight edges, coupling edge sets

\mathcal{E}_{\mathrel{\vbox{\hbox{\scalebox{0.75}{

\rightrightarrows

}}}}}(\cdot,\cdot)

have type coupling , and the edges joining

u^{h,\ell}

to its head have type head .

	Corpus-mixture dataset	Fine-tuning dataset
Number of models	$3{,}000$	$4{,}000$
Parameters per model	$91$ K– $150$ K (mean $121$ K)	$3.8$ M– $5.1$ M (mean $4.4$ M)
Vocabulary (characters)	$64$	$65$
Embedding dim. $d$	$64$	$320$
Blocks $L$	$\{3,5\}$	$6$
Heads $H$	$\{1,2,4\}$	$\{1,2,4\}$

Appendix

Table 7: Summary of the two transformer-population datasets.

Method	Input	Architecture
GNN (all constructions)	graph	$3$ bidirectional GINE layers, hidden size $64$ , no dropout, trainable Fourier weight encoder, MLP readout
NFN-T [ 34 ]	weight tensors	one equivariant layer ( $1$ channel) + invariant pooling + MLP; fixed architecture only
DeepSets	set of $N$ scalars	weight embedding ( $128$ ), $3$ equivariant layers, mean-pool, MLP $[128]$ , dropout $0.05$
MLP	vector in $\mathbb{R}^{N}$	hidden layers $[512,256]$ , dropout $0.1$ (reduced on the fine-tuning dataset, see text)
MLP+Aug	vector in $\mathbb{R}^{N}$	same as MLP, with symmetry-group augmentation of the training weights

Appendix

Table 8: Metanetwork architectures.

N

denotes the number of parameters of the input transformer (padded to the largest model in the population for the flat baselines). All baselines are trained with the same optimizer and protocol as the graph metanetworks ( Section C.3 ).

Panel of Figure 3 — Figure 3: Full construction of a single-block transformer, containing a 2-head attention mechanism. We omit the coupling between the first and last layers, as well as the normalization nodes and edges, to keep the plot intuitive and readable.

Panel of Figure 4 — Figure 4: Full construction of a single-block transformer, containing a 2-head attention mechanism. Attention construction is invariant to permutations only. We omit the coupling between the first and last layers, as well as the normalization nodes and edges, to keep the plot intuitive and readable.

Weight-Space Learning

See all 17 papers

Interfaces and Modules: A Scalable Representation for Weight-Space Learning across Architectures
Oct 15, 2026Michał Kozyra, Michael BronsteinWeight-Space Learning Weight-Space Analysis
University of Oxford · AITHYRA
Anatomical Heterogeneity in Transformer Language Models
Mar 19, 2026Tomasz WietrzykowskiTransformer Efficient Language Model Training
Independent Researcher, Wroclaw, Poland
W2T: LoRA Weights Already Know What They Can Do
Mar 16, 2026Xiaolong Han, Ferrante Neri, Zijian Jiang +4Weight-Space Learning Low-Rank Adaptation
University of Surrey · Stanford University · University of Notre Dame

Oct 15, 2026J/K move · Enter open · S save

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

Michał Kozyra, Michael Bronstein

University of Oxford · AITHYRA

Weight-Space Learning Weight-Space Analysis Equivariant Neural Networks Permutation-Equivariant Neural Networks Graph Representation Learning

Open paper AskPDF arXiv

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.

Abstract

Weight-Space Learning

Explore similar work

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

Anatomical Heterogeneity in Transformer Language Models

W2T: LoRA Weights Already Know What They Can Do