Weight-Space Mixture-of-Experts for Implicit Neural Representation Classification

AI summary

Overview

• Research area: Computer vision / representation learning — specifically image classification performed in the weight space of Implicit Neural Representations (INRs), combined with Mixture-of-Experts architectures, meta-learning, and weight-space explainability.
• Technical level: Advanced. The paper assumes familiarity with implicit neural representations (SIREN), MAML/Meta-SGD meta-learning, Transformer tokenization, sparse MoE routing, and gradient-based attribution.
• Scope: The paper proposes a hierarchical Mixture-of-Experts (HMoE) Transformer for classifying INR weights, trains it end-to-end with a meta-learned SIREN, and adds two weight-space explainability tools (gradient attribution and structured pruning), evaluated on MNIST, Fashion-MNIST, CIFAR-10, Imagenette, and ImageNet-1K.

What This Paper Is About

Normally, an image classifier looks at pixels. An alternative is to first fit a small coordinate-based network (an INR) to each image, then classify the weights of that network instead of the image itself. This weight-space view is attractive but hard: INR weights are high-dimensional, contain permutation symmetries, and it is unclear which parameters actually carry class information. The paper's goal is to build a classifier that exploits the layered structure of INR weights, and to develop tools that reveal where discriminative information lives inside those weights.

Key Contributions

1. A hierarchical Mixture-of-Experts (HMoE) Transformer for INR weight space. Each standard Transformer feed-forward block is replaced by a two-stage sparse MoE: a Layer-wise MoE that routes whole INR-layer token groups, followed by a Token-wise MoE that routes individual tokens conditioned on the layer-level routing signal.
2. State-of-the-art weight-space classification results across low-resolution benchmarks (MNIST, Fashion-MNIST, CIFAR-10) and high-resolution benchmarks (Imagenette, ImageNet-1K), including a new CIFAR-10 accuracy of 69.11% when scaled to 20 blocks with augmentations, and 26.73% on ImageNet-1K.
3. Weight-space explainability methods. Gradient-Weighted Class Activation Graph (Grad-CAG), a weight-space analogue of Grad-CAM that scores each INR weight by its first-order sensitivity to a class logit, plus a structured pruning procedure for producing spatial attribution maps.
4. An ablation study isolating the effects of layer-wise routing, expert count, active-expert count (TopK), task-loss weighting, inner-loop step count, and SIREN depth.

Main Findings

• Low-resolution benchmarks: HMoE-MWT achieves 99.06% on MNIST, 90.72% on Fashion-MNIST, and 65.01% on CIFAR-10, surpassing the previously reported MWT (98.33%, 89.41%, 56.90%) and all listed INR- and graph-based baselines (DWS, NFN, ScaleGMN, NG-GNN, Inr2Vec, MLP). Even the non-meta-learned HMoE-WT reaches 97.87% / 87.57% / 55.44%.
• Architecture matters more than width. Unlike the prior MWT work, the authors observe no accuracy gain from increasing SIREN width: HMoE-MWT-L (width 256) scored 98.93% / 90.30% / 64.40%, slightly below the width-128 HMoE-MWT. HMoE-MWT also surpasses the MWT-L baseline trained for the same extended schedule.
• CIFAR-10 scale-up: A 20-transformer-block HMoE model with augmentations reached 69.11%, above the 64.7% reported by the original MWT authors for their large 20-block model trained for 40 epochs with augmentations.
• Imagenette: HMoE-MWT attains 62.52% accuracy / 21.47 dB PSNR, and HMoE-MWT-L reaches 62.96% / 23.01 dB, both above MWT (56.78% / 21.14 dB) and MWT-L (60.62% / 22.31 dB). The non-meta-learned HMoE-WT (55.39% / 22.49 dB) already matches or exceeds MWT.
• Parameter-matched comparison confirms the gain is architectural. Compressing HMoE-MWT to 1.1M classifier parameters (8 blocks, FFN ratio 0.25) still yields 61.47% versus MWT's 56.78% at the same 1.1M budget; expanding MWT to ~11M parameters yields 59.92%, below HMoE-MWT's 62.52%.
• Sparse routing keeps compute comparable. With TopK k=1 and four experts per stage, only one expert is active per token, so per-token FLOPs stay similar to dense models while total capacity grows. Reported FLOPs: HMoE-MWT 8.8G versus MWT 2.3G and MWT-L 19G on Imagenette.
• ImageNet-1K: HMoE-MWT-L reaches 26.73% accuracy with 22.06 dB PSNR, compared with 24.11% / 21.78 dB for MWT-L.
• Weight attribution is functionally meaningful. Pruning the highest-attribution weights sharply degrades reconstruction, while pruning the lowest-attribution weights preserves PSNR longer than magnitude-based pruning beyond roughly 40% of weights removed. This suggests class-relevant information sits in subcircuits spanning diverse weight magnitudes.
• Spatial localisation. On a manually annotated set of 200 segmentation masks (20 per class) from the Imagenette validation set, attribution maps scored 0.35 on the pointing game, versus 0.22 for magnitude-based pruning and 0.15 for random pruning.
• Failure modes are exposed. Attribution maps reveal a dataset bias in Imagenette: the tench class is spuriously associated with the fishing net, a frequently co-occurring object.
• Manifold structure separates the variants. UMAP embeddings of Grad-CAG-weighted parameters show HMoE-MWT forming smooth, overlapping manifolds that generalise from train to validation splits, while HMoE-WT produces fragmented, class-aligned clusters that shift between splits — consistent with its lower accuracy.
• Layer-wise routing is essential. Removing the layer-wise MoE (token-wise MoE only) drops CIFAR-10 accuracy from 65.01% to 62.57% (TopK 1), 62.54% (TopK 2), or 62.21% (TopK 4). Adding more experts slightly improved PSNR but not accuracy (e.g., 1 layer / 8 token experts: 64.53%, 32.27 dB, 17M parameters, versus 65.01% / 31.71 dB / 11M for the default).
• Task-loss weighting has a clear optimum. On an 80/20 CIFAR-10 split, w_task = 0.01 gave 63.14% accuracy, versus 62.63% at 0.001, 58.54% at 0.1, and 55.34% at 1.
• More inner-loop steps help. k = 1, 2, 4, 6 steps gave 57.84%, 60.23%, 62.22%, and 63.14% accuracy respectively on the 80/20 split.

Methodology in Plain English

Each training image is represented by a SIREN (a sinusoidal-activation MLP) that maps pixel coordinates to RGB values. Instead of fitting a fresh network per image from scratch, the authors learn a shared starting point (an initialization θ) plus a per-parameter learning-rate schedule α. For every image, a small inner loop of gradient steps adapts θ into an image-specific parameter set φ using a reconstruction loss.

Crucially, the classification loss is also backpropagated into this adaptation process, so the shared initialization is shaped not only to reconstruct images well but also to be easy to classify. The classifier itself does not see pixels — it sees the hidden-layer weights of the adapted SIREN, with biases folded into the weight matrices so each output neuron becomes one token. A learned positional encoding is added, and the classifier operates on the scaled difference λ(φ − θ) with λ = 500, which emphasises what changed during adaptation.

Inside the classifier, each feed-forward block is replaced by a two-stage MoE. Stage 1 averages all tokens belonging to the same INR layer into a summary vector, routes that summary through a softmax gate, picks the top-k experts, and sends the entire layer's tokens to those experts — preserving layer-level structure. Stage 2 routes each individual token, using a gating score that combines the token's own features with the layer-level routing vector projected into token-expert space. A load-balancing loss encourages all experts to be used. Defaults are 10 Transformer blocks, SIREN with four hidden layers, two-layer FFN experts, single-linear gating layers, TopK k=1, four experts per stage, 40 epochs, and loss weights of 1 (reconstruction), 0.01 (classification), and 0.1 (balancing), trained on an NVIDIA RTX 4090.

For explainability, Grad-CAG assigns each INR weight an importance score equal to the absolute value of the weight times the gradient of the target class logit with respect to that weight. These scores drive pruning experiments (prune from least to most important, most to least important, randomly, or by magnitude) and are converted into spatial attribution maps by pruning 60% of weights, reconstructing the image, inverting the difference from the original reconstruction, downsampling to 32×32, applying a Gaussian filter with σ = 2, upsampling back, and applying a nonlinear intensity remapping.

Why This Matters

Impact on research. Weight-space learning has been held back by two problems: classifiers that ignore the internal structure of INRs, and a near-total absence of interpretability tools for weight-space models. This work addresses both. The findings that layer-aware routing consistently beats token-only routing, and that classifier architecture matters more than INR width, are concrete design lessons for anyone working on neural functionals or INR-based learning. Grad-CAG is presented as the first interpretability tool for weight-space classification.

Real-world applications (as motivated by the paper):

• Privacy-preserving and federated learning — models operate on network parameters rather than explicit signal samples, which the authors frame as an advantage where sharing raw data is restricted.
• Biomedical applications — the same privacy argument is extended to medical settings where raw patient data cannot be freely distributed.
• Multimodal signal processing — images, audio, and other signals can all be mapped into INR parameters and processed by one consistent framework; INRs can also jointly represent multiple quantities.
• Image representation, synthesis, and compression — INRs originated in these areas, and the pruning analyses here quantify how aggressively INR weights can be removed while preserving reconstruction quality.

Industry relevance. The sparse top-k routing means only one expert fires per token, so the reported FLOPs per token stay comparable to dense baselines even though total parameter capacity grows. That is an attractive trade-off for deployment, where capacity can be bought without proportionally increasing inference cost. The paper also directly engages with efficiency limits: MoE models need longer training and higher compute, and encoding each image into an INR remains a per-sample cost despite meta-learning amortizing it.

Future Directions

• Extending beyond image SIRENs. The authors state that all current experiments use image-based SIRENs and that testing the HMoE framework on other INR architectures and domains — 3D NeRFs are named explicitly — is left for future work.
• Reducing the cost of INR encoding. Training one model per image is described as a remaining computational bottleneck, even with meta-learning.
• Closing the gap to pixel-based models. Despite the gains within weight space, performance still lags conventional pixel-based models on large-scale datasets such as ImageNet-1K, where 26.73% accuracy was reached.
• Deepening weight-space interpretability. The paper characterises its explainability analysis as a preliminary first step covering only weight importance and pruning, leaving many aspects of weight-space interpretability unexplored.
• Reducing MoE training cost. MoE models were found to require longer training and higher computational cost than standard Transformers, and the authors found longer schedules beneficial (40 epochs by default), which raises the question of whether more efficient training schemes are possible.

Target Audience

This paper is for researchers working on implicit neural representations, neural functionals, and weight-space learning who want a stronger classifier backbone and a first look at interpreting INR weights. It also suits practitioners of Mixture-of-Experts who are interested in non-standard token structures — here, tokens are organised by network layer and by originating neuron, which makes layer-aware routing natural. Finally, it is relevant to explainability researchers, since Grad-CAG extends gradient attribution from pixel space into a domain where pixel-space attribution methods cannot be applied at all. Readers should already be comfortable with SIREN, MAML/Meta-SGD meta-learning, Transformer tokenization, and sparse MoE routing; the presentation itself is standard conference-style with equations, tables, and figures.

Authors’ abstract

Implicit Neural Representations (INRs) encode signals as the weights of a coordinate-based neural network and have recently been proposed as an alternative domain for downstream learning. While promising, classification directly in weight space remains challenging due to the high dimensionality and complex structure of INR parameters. Furthermore, the way discriminative information is distributed across INR weights remains poorly understood. We propose a hierarchical Mixture-of-Experts (HMoE) Transformer that processes INR weights using conditional computation aligned with the structure of the underlying implicit network. Coupled with a meta-learning framework that shapes INR parameters for downstream tasks, our model achieves state-of-the-art accuracy across standard benchmarks, ranging from low-resolution datasets to high-resolution ImageNet-1K. To gain insight into how INRs encode discriminative information, we develop weight-space attribution and pruning methods that identify parameters most relevant for classification. These analyses reveal how class-specific structure emerges within INR layers and support the suitability of MoE architectures for weight-space learning. Our approach advances both the performance and interpretability of weight-space classifiers.