Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper introduces a novel typed graph representation, termed interfaces and modules, to enable scalable weight-space learning across diverse neural network architectures by explicitly encoding coordinate-permutation symmetries. This representation allows for efficient global attention over a small number of invariant interface summaries, overcoming the limitations of architecture-specific encoders and costly neuron graph attention. The Universal Weight Encoder (UWE) leverages this structure to achieve competitive performance on existing benchmarks and significantly outperform baselines on a large zoo of transformer architectures, even for unseen and larger models, with reduced memory footprint.
A single weight-space encoder now masters 100 transformer architectures, outperforming specialized models and handling unseen, larger networks with less memory.
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.