Structure Is Not Enough: Leveraging Behavior for Neural Network Weight Reconstruction

TL;DR

This work shows that purely structural reconstruction in weight-space autoencoders is insufficient to faithfully recover high-performing neural networks. By introducing a query-based behavioral loss that compares model outputs on selected inputs, and integrating it with a structural loss in a composite objective, the authors demonstrate strong synergistic gains across reconstructive and generative tasks. Gradient analysis reveals that the behavioral loss aligns weight perturbations with input-output sensitivities, offering functional guidance missing from purely structural signals. Empirically, combining structure and behavior on three model zoos yields reconstructions and generated weights that approach the performance of original models, suggesting practical potential for weight-space data generation and analysis with a data-efficient, self-supervised framework.

Abstract

The weights of neural networks (NNs) have recently gained prominence as a new data modality in machine learning, with applications ranging from accuracy and hyperparameter prediction to representation learning or weight generation. One approach to leverage NN weights involves training autoencoders (AEs), using contrastive and reconstruction losses. This allows such models to be applied to a wide variety of downstream tasks, and they demonstrate strong predictive performance and low reconstruction error. However, despite the low reconstruction error, these AEs reconstruct NN models with deteriorated performance compared to the original ones, limiting their usability with regard to model weight generation. In this paper, we identify a limitation of weight-space AEs, specifically highlighting that a structural loss, that uses the Euclidean distance between original and reconstructed weights, fails to capture some features critical for reconstructing high-performing models. We analyze the addition of a behavioral loss for training AEs in weight space, where we compare the output of the reconstructed model with that of the original one, given some common input. We show a strong synergy between structural and behavioral signals, leading to increased performance in all downstream tasks evaluated, in particular NN weights reconstruction and generation.

Figures (10)

• Figure 1: Illustration of our behavioral loss. Previous works have used a structural loss, comparing NN weights with their reconstruction using the Euclidean distance. In that setting however, they failed to reconstruct models that perform as a well as the original. We propose combining it with a behavioral loss, where we use queries to compare the outputs of the NN and its reconstruction. Since this loss does not attempt to make the reconstructed model accurate with regard to a ground truth, rather to match its behavior with that of the original, the queries do not need to be labeled.
• Figure 2: Evaluation of the reconstructive downstream tasks, pairwise between models from the model zoo and their reconstructions, depending on the losses used to train the hyper-representation AE. Each column represents one of our model zoos. Top: distribution of pairwise structural $L^2$ distances. Bottom: distribution of the pairwise behavioral similarities, measured with model agreement. On the structural side, we see that using the structural loss is sufficient to concentrate most pairwise distances around some low value. On the behavioral side, we see that using the behavioral loss only yields the worst performance, even compared to the structural loss. Using both the structural and behavioral losses is necessary to achieve high levels of agreement. Most models show high levels of agreements, but since a few show low levels of agreement the standard deviation shown in Table \ref{['tab:reconstructive_dstk']} can be high.
• Figure 3: Evaluation of the reconstructive and generative downstream tasks, shown as distributions of the test accuracy of different models, depending on whether they are part of the original model zoo, reconstructions of models from that model zoo, or generated models. Each column represents one of our model zoos, while the row show what loss has been used to train the specific hyper-representation model. The top row shows results for our baseline, that uses the contrastive $\mathcal{L}_C$ and structural $\mathcal{L}_S$ losses. The bottom row represents our hyper-representation AEs, which in addition are also trained with a behavioral loss $\mathcal{L}_S$. We note that for the baseline, neither the reconstructed nor the generated models can match the performance of the original models. On the other hand, when adding a behavioral element to the loss, they match the performance of the most accurate models from the original zoo.
• Figure 4: Distribution of the model test accuracies in our three model zoos, by training epoch.
• Figure 5: Comparison of the performance of the discriminative downstream tasks by training epoch of SANE. We first note that both models that use $\mathcal{L}_S$ and $\mathcal{L}_B$, show a more stable increase in performance with each epoch, and mostly outperform other models. Then, comparing the model that uses $\mathcal{L}_S \oplus \mathcal{L}_B$ with the one that uses $\mathcal{L}_C \oplus \mathcal{L}_S \oplus \mathcal{L}_B$, we qualitatively note that in most cases, performance grows faster with the number of epochs.