Weight-Space Mixture-of-Experts for Implicit Neural Representation Classification¶
Conference: ECCV 2026
Paper: ECCV Official
Code: https://github.com/stasiek-j/HMoE-MWT
Area: Interpretability
Keywords: Implicit Neural Representations (INRs), Weight-Space Learning, Mixture-of-Experts (MoE), Hierarchical Routing, Weight-Space Explainability (Grad-CAG)
TL;DR¶
Addressing the challenges of high parameter dimensionality and layer-wise functional heterogeneity in classifying implicit neural representation (INR) weights, this paper proposes a hierarchical Mixture-of-Experts (HMoE) Transformer coupled with meta-learning that achieves new state-of-the-art results across standard benchmarks and introduces Grad-CAG, a first-order weight-space attribution and pruning method to uncover discriminative subcircuits.
Background & Motivation¶
Implicit Neural Representations (INRs) parameterize signals via coordinate-based multilayer perceptrons (MLPs), mapping low-dimensional coordinates to target signal intensities. By replacing discrete pixel grids with network parameters, INRs offer a unified framework capable of handling diverse modalities—from images and audio to 3D scenes—within a consistent weight-space representation. This continuous formulation is particularly advantageous for privacy-preserving machine learning, such as federated learning or medical image analysis, where sharing network weights is preferred over exposing raw data.
However, operating directly in INR weight space remains challenging. MLPs possess pervasive permutation and scaling symmetries that introduce non-trivial representation ambiguities, high parameter counts necessitate aggressive dimensional compression, and fitting separate INRs per instance incurs prohibitive computational overhead. While the Meta Weight Transformer (MWT) alleviated initialization divergence and structured weight manifolds via meta-learning (MAML/Meta-SGD), existing weight-space classifiers treat all weight neurons as interchangeable tokens and pass them through identical feed-forward networks (FFNs). This homogeneous treatment fundamentally ignores the core architectural inductive bias of INRs—namely that distinct network layers encode disparate frequency bands and that individual neurons specialize into modular functional roles.
Mixture-of-Experts (MoE) architectures naturally introduce conditional computation and functional modularity by routing inputs to specialized sub-networks, matching the multi-scale frequency decomposition and subnetwork modularity inherent to INR parameter spaces. Core idea: construct a hierarchical Mixture-of-Experts (HMoE) Transformer aligned with INR network topology—combining layer-level coarse routing with cross-level biased token routing—and establish the first weight-space attribution framework (Grad-CAG) to systematically explain decisions and locate task-relevant subcircuits.
Method¶
Overall Architecture¶
The framework unites an end-to-end meta-learning pipeline with a hierarchical weight-space classification Transformer. In the meta-learning stage, a shared initialization \(\theta\) and a per-step coordinate learning rate schedule \(\alpha\) are maintained. For any given input sample, the inner loop adapts \(\theta\) across \(k\) gradient steps using mean squared error reconstruction loss, producing sample-specific INR parameters \(\phi\). The classifier constructs neuron tokens from the hidden layers by concatenating incoming weights with their folded biases, forming tokens of dimension \(D = n + 1\) and a sequence of length \(M = nL\). The network processes scaled parameter residuals \(\lambda(\phi - \theta)\) (with scaling factor \(\lambda = 500\)), routing tokens through 10 stacked HMoE blocks before average-pooling into class logits.
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["Input Image & Coordinate Grid"] --> B["Meta-Learning Inner Loop Adaptation<br/>k-step gradient updates produce INR weights ϕ"]
B --> C["Weight Residual Tokenization<br/>λ(ϕ - θ) grouped into layer-wise blocks"]
C --> D["Layer-wise MoE Routing<br/>Mean-pooled layer token guides Top-1 layer expert dispatch"]
D --> E["Cross-Stage Gate Projection<br/>Linear mapping projects layer gate to bias token gating"]
E --> F["Fine-grained Token-wise MoE<br/>Top-1 expert sparse conditional computation"]
F --> G["Joint Classification & Load Balancing Loss<br/>L_cls backpropagates to meta-INR + L_balance prevents collapse"]
G --> H["Weight-Space Attribution Grad-CAG<br/>First-order importance scores guide subcircuit pruning"]
Key Designs¶
1. Layer-wise MoE Routing: Decoupling coarse layer frequency structure from uniform token transformations
Dense Transformers treat all INR tokens uniformly, failing to capture the layer-wise frequency stratification where early layers capture low frequencies and deeper layers synthesize fine details. To preserve this inductive bias, the model groups the full sequence \(X \in \mathbb{R}^{M \times D}\) by its originating INR layer into \((X_1, \dots, X_L)\), where each \(X_\ell \in \mathbb{R}^{n \times D}\) aggregates all \(n\) neuron tokens of layer \(\ell\). The module computes a summary representation for the entire layer via mean pooling \(v_\ell = \frac{1}{n} \sum_{i=1}^n X_\ell[i]\) and evaluates the layer routing vector using linear projection \(W_g\) followed by softmax: $\(g_\ell = \text{softmax}(W_g v_\ell)\)$ A sparse selection operator \(\text{TopK}(g_\ell, k)\) then selects the top-\(k\) layer experts \(E_j\). All tokens belonging to layer \(\ell\) are processed concurrently by these chosen experts and recombined according to their gating probabilities: \(X_\ell' = \sum_{j \in \text{TopK}(g_\ell, k)} g_{\ell, j} E_j(X_\ell)\). Re-concatenating layer outputs in order maintains the macro structural topology while enabling functional specialization across layers.
2. Cross-Stage Gate Projection: Propagating macro layer context to guide fine-grained token expert allocation
While layer-wise routing captures global structural roles, individual neurons within the same layer exhibit distinct functional specializations. Unconstrained token-level routing would disconnect individual neurons from their macro layer context. To enforce cross-scale coherence, the architecture introduces a learnable projection matrix \(B \in \mathbb{R}^{E_L \times E_T}\) that maps the layer routing probabilities \(g_\ell\) from Stage 1 into the token expert space. For any token \(x \in X_\ell\) belonging to layer \(\ell\), the token-level routing score vector \(s\) is computed as: $\(s = \text{softmax}(W_s x + g_\ell B)\)$ Each token independently activates its top-\(k\) token experts according to \(s\): \(x' = \sum_{j \in \text{TopK}(s, k)} s_j E_j(x)\). This hierarchical coupling conditions fine-grained token routing on macro layer-level routing, aligning localized neuron specialization with global architectural context.
3. Meta-Learning Task-Driven Weight Alignment: End-to-end backpropagation to structure discriminative weight manifolds
Rather than treating INR fitting as an independent, isolated preprocessing phase, the framework trains the meta-initialization \(\theta\) and the per-parameter inner-loop learning rates \(\alpha\) jointly with the classifier. In the inner loop, parameters adapt via pixel reconstruction loss. In the outer loop, the downstream classification cross-entropy loss \(L_{cls}\) backpropagates through the entire adaptation trajectory into \(\theta\) and \(\alpha\). This direct coupling guides the meta-INR toward weight trajectories that are not only high-fidelity for coordinate reconstruction but also cleanly separated and easily discriminable for downstream classification.
4. Weight-Space Gradient-Weighted Class Activation Graph (Grad-CAG): Quantifying first-order sensitivity for subcircuit discovery
Because the classifier operates exclusively on network parameters, standard pixel-space attribution methods (e.g., Grad-CAM) cannot be directly applied. Grad-CAG defines a first-order importance score for each scalar INR weight \(\phi_j\) with respect to target class \(i\): $\(S_j^{(i)} = \left| \phi_j \frac{\partial c_i(\phi)}{\partial \phi_j} \right|\)$ This metric quantifies the local sensitivity of the class prediction to fractional perturbations in \(\phi_j\). Flattened attribution vectors projected via UMAP reveal the underlying geometry of the weight manifold. Furthermore, pruning weights guided by \(S_j^{(i)}\) isolates task-specific subcircuits: rendering pruned INRs and computing spatial residuals against the original reconstruction enables spatial localization of class-discriminative evidence directly from weight space.
Loss & Training¶
The overall training objective combines the coordinate mean squared reconstruction loss \(L_{rec}\), the classification cross-entropy loss \(L_{cls}\), and a dual-stage load balancing loss \(L_{balance}\) to prevent expert collapse. The balancing loss computes the inner product between expert assignment frequency and mean gating probability across the batch: $\(\mathcal{L}_{\text{balance}}^{\text{layer}} = E_L \sum_{e=1}^{E_L} s_e P_e, \quad \mathcal{L}_{\text{balance}}^{\text{token}} = E_T \sum_{j=1}^{E_T} \tilde{s}_j \tilde{P}_j\)$ The module averages both auxiliary losses and optimizes the composite objective: $\(\mathcal{L}_{\text{total}} = \mathcal{L}_{rec} + w_{task} \mathcal{L}_{cls} + w_{balance} \mathcal{L}_{balance}\)$ By default, \(w_{task} = 0.01\) and \(w_{balance} = 0.1\). Both layer and token stages configure \(E_L = 4, E_T = 4\) with Top-1 routing (\(k=1\)), ensuring that per-token inference FLOPs remain identical to standard dense Transformers. Models are trained with AdamW (learning rate 0.0001, weight decay 0.0001) for 40 epochs on an NVIDIA RTX 4090 GPU.
Key Experimental Results¶
Main Results¶
The model is benchmarked across standard low-resolution datasets (MNIST, Fashion-MNIST, CIFAR-10) and high-resolution datasets (Imagenette, ImageNet-1K), consistently outperforming prior weight-space and graph-based classifiers.
| Dataset | Metric | HMoE-MWT (Ours) | MWT-L* (Baseline) | Gain |
|---|---|---|---|---|
| MNIST | Accuracy (%) | 99.06 ± 0.15 | 98.91 ± 0.11 | +0.15% |
| Fashion-MNIST | Accuracy (%) | 90.72 ± 0.25 | 90.35 ± 0.20 | +0.37% |
| CIFAR-10 | Accuracy (%) | 65.01 ± 0.63 | 58.90 ± 0.35 | +6.11% |
| Imagenette | Accuracy (%) | 62.52 ± 0.60 | 60.62 ± 0.37 (MWT-L) | +1.90% |
| ImageNet-1K | Top-1 Accuracy (%) | 26.73 (HMoE-MWT-L) | 24.11 (MWT-L) | +2.62% |
Under strict parameter-matched constraints on Imagenette: at the 1.1M parameter budget, [email protected] achieves 61.47 ± 0.32% accuracy (1.9G FLOPs), outperforming dense MWT's 56.78 ± 0.27% (2.3G FLOPs); at the 11M parameter budget, HMoE-MWT reaches 62.52 ± 0.60% (8.8G FLOPs), outperforming scaled dense MWT-PM@11M's 59.92 ± 0.35% (15.6G FLOPs). Scaling HMoE-MWT to 20 blocks with data augmentations on CIFAR-10 establishes a new SOTA accuracy of 69.11%.
Ablation Study¶
The architectural ablation on CIFAR-10 systematically evaluates routing stages and expert counts within the 10-block classifier:
| Routing Configuration | TopK (Layer) | # Layer Exp | TopK (Token) | # Token Exp | Acc [%] | PSNR [dB] | Params | Note |
|---|---|---|---|---|---|---|---|---|
| Layer MoE only | 1 | 4 | - | - | 63.21 | 33.86 | 11M | Lacks fine token-level adaptation |
| Token MoE only | - | - | 1 | 4 | 62.57 | 30.70 | 11M | Loses global layer-level context |
| Token MoE only (\(k=2\)) | - | - | 2 | 4 | 62.54 | 31.14 | 11M | Activating more experts provides no gain |
| Token MoE only (\(k=4\)) | - | - | 4 | 4 | 62.21 | 30.68 | 11M | Dense computation causes routing blur |
| Full HMoE (Default) | 1 | 4 | 1 | 4 | 65.01 | 31.71 | 11M | Optimal two-stage cooperation |
| Expanded Layer Capacity | 1 | 8 | 1 | 4 | 64.03 | 31.86 | 17M | Extra layer capacity does not boost acc |
| Expanded Token Capacity | 1 | 4 | 1 | 8 | 64.53 | 32.27 | 17M | Slight reconstruction gain only |
| Dual-stage Expanded Capacity | 1 | 8 | 1 | 8 | 63.64 | 30.61 | 22M | Overfitting leads to minor degradation |
Key Findings¶
- Hierarchical layer-token routing is critical: Removing layer-wise routing and relying solely on token-wise MoE causes accuracy to drop from 65.01% to 62.57%, proving that capturing macro layer structure is essential for weight-space representation learning.
- Sparse routing surpasses dense computation: Increasing \(k\) in token routing from 1 to 2 or 4 reduces accuracy to 62.54% and 62.21% while increasing computational cost, confirming that sparse \(k=1\) routing provides sufficient functional specialization.
- Grad-CAG pruning preserves object geometry: Pruning up to 60% of the lowest-attribution weights preserves coarse object structure and yields a Pointing Game localization score of 0.35, substantially outperforming magnitude-based pruning (0.22) and random pruning (0.15).
- Meta-learning prevents manifold fragmentation: UMAP visualizations show that without meta-learning, HMoE-WT produces fragmented, class-isolated clusters with validation drift, whereas HMoE-MWT yields continuous, shared functional subcircuits that generalize reliably across splits.
Highlights & Insights¶
- Natural alignment between network topology and conditional computation: The two-stage MoE design mirrors the physical organization of coordinate MLPs—mean-pooled layer routing preserves frequency decomposition, while cross-stage gate projection biases fine-grained neuron specialization.
- First dedicated weight-space explainability suite: Grad-CAG introduces a rigorous first-order attribution formulation for parameters, opening a new path for auditing black-box weight-space models and revealing dataset co-occurrence biases (such as the tench-fishing net spurious correlation).
- Superior compute-parameter pareto efficiency: By setting Top-1 routing, the model increases total parameter capacity to 11M without increasing per-token active FLOPs, outperforming both parameter-matched and compute-matched dense baselines.
Limitations & Future Work¶
- Per-sample fitting latency: While meta-learning amortizes fitting time via \(k\)-step inner loops, requiring an independent optimization stage for each test sample remains a deployment bottleneck compared to feed-forward vision backbones.
- Evaluation restricted to 2D image SIRENs: Experiments focus exclusively on sinusoidal coordinate MLPs on 2D images; extending hierarchical MoE routing to 3D NeRFs, 3D Gaussian Splatting, or audio representations remains unexplored.
- First-order local attribution limitations: Grad-CAG relies solely on local first-order derivatives and parameter magnitudes, omitting higher-order parameter interactions and path-integrated gradient accumulations.
Related Work & Insights¶
- vs ScaleGMN / DWS-Net / NFN: Early weight-space architectures enforce strict permutation and scaling equivariance through specialized tensor operations, but scale poorly to deep, wide networks. This paper uses meta-learning to canonicalize weight distributions and leverages scalable sparse MoE routing to achieve higher accuracy on larger benchmarks.
- vs Meta Weight Transformer (MWT): MWT introduced joint meta-INR and dense Transformer training. This work replaces the uniform dense FFN with a two-stage hierarchical MoE, producing a +6.11% boost on CIFAR-10, scaling successfully to full ImageNet-1K, and providing weight-space attribution tools.
Rating¶
- Novelty: ⭐⭐⭐⭐ [Introduces hierarchical MoE conditional computation tailored to INR layer topology and presents the first weight-space gradient attribution method]
- Experimental Thoroughness: ⭐⭐⭐⭐⭐ [Extensive evaluations spanning low-resolution to full ImageNet-1K, parameter-matched comparisons, routing ablations, and pointing-game localization audits]
- Writing Quality: ⭐⭐⭐⭐⭐ [Clear structural narrative, sound theoretical justification, and detailed supplementary experimental protocols]
- Value: ⭐⭐⭐⭐ [Provides an effective blueprint for learning on neural representations and interpreting learned subcircuits in weight space]