Research map / Architectures and efficiency / RNNs and normalization

Universal Neural Functionals

NeurIPS 2024 · Architectures and efficiency › RNNs and normalization · approach: Symmetry and equivariance in loss landscapes · 2 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes an algorithm that automatically constructs permutation-equivariant linear layers (universal neural functionals, UNFs) for arbitrary weight spaces and applies them to RNN generalization prediction and learned optimizers.

Problem

Processing weight-space features (weights, gradients) of general neural architectures with recurrence, residuals, or attention is hard because existing permutation-equivariant weight-space models only handle simple MLPs and CNNs.

Approach

Automatically enumerates valid partitions of tensor indices that permute simultaneously to build a basis for the maximally expressive permutation-equivariant linear layer over any collection of tensors, which is then stacked into deep equivariant models usable as learned optimizers or generalization predictors.

Key results

• On Tiny RNN Zoo generalization prediction, UNF achieves test Kendall's tau 0.8968 vs 0.8839 for StatNN.
• UNF-based learned optimizers achieve the lowest training loss across the four tasks, typically learning fastest; Deep Sets and NFN can be comparable in some settings.
• UNF outperforms NFN in the CNN experiment, attributed to treating spatial kernel dimensions as permutable, which yields fewer parameters.

Measured effects reported in the paper

Surprising findings

• Assuming spatial kernel dimensions are permutable is technically incorrect, yet the stronger assumption performs better in the CNN learned optimizer experiment.

Limitations

• Computational tractability degrades for higher-rank tensors as the number of basis terms grows rapidly.
• Not yet applicable to heterogeneous weight-space inputs (a single UNF across architectures).
• Learned optimizer experiments are small-scale and generalization to new tasks is untested.

Experimental setup

Datasets: Tiny RNN Zoo FashionMNIST CIFAR-10 LM1B

Benchmarks: Tiny RNN Zoo FashionMNIST CIFAR-10 LM1B

Metrics: Kendall's tau training loss

Baselines: Deep Sets learned optimizer NFN_NP SGDM StatNN

Builds on: Neural Functional Networks (NFN) Deep Sets persistent evolutionary strategies StatNN

Closest papers by meaning

• On Universality Classes of Equivariant Networks (NeurIPS 2025 · similarity 0.87)
• Permutation Equivariant Neural Networks for Symmetric Tensors (ICML 2025 · similarity 0.86)
• Monomial Matrix Group Equivariant Neural Functional Networks (NeurIPS 2024 · similarity 0.86)
• Equivariant Polynomial Functional Networks (ICML 2025 · similarity 0.85)
• On the hardness of learning under symmetries (ICLR 2024 · similarity 0.85)
• Learning symmetries via weight-sharing with doubly stochastic tensors (NeurIPS 2024 · similarity 0.85)