Data Augmentation Techniques to Reverse-Engineer Neural Network Weights from Input-Output Queries

2025-11-25

AI Summary

In teacher-student paradigms, weight inversion fails when model parameters vastly exceed training samples—causing student networks to overfit query inputs rather than align with teacher parameters. To address this, we propose a hidden-layer representation–aware data augmentation strategy. Unlike conventional augmentations (e.g., rotation, flipping, or noise injection), our method actively elicits diverse internal representations from the teacher network, thereby enhancing the student’s capacity for functional approximation and parametric alignment with the teacher. Experiments demonstrate successful reconstruction of teacher weights in networks whose parameter count is up to 100× the number of training data points—surpassing prior scalability limits in weight inversion. This work establishes, for the first time, the feasibility and scalability of input-output query–driven weight recovery in large-scale neural networks.

Technology Category

• Machine Learning: Adversarial Learning & Robustness
• Search and Optimization: Learning to Search
• Computer Vision: Representation Learning for Vision

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

Abstract

Network weights can be reverse-engineered given enough informative samples of a network's input-output function. In a teacher-student setup, this translates into collecting a dataset of the teacher mapping -- querying the teacher -- and fitting a student to imitate such mapping. A sensible choice of queries is the dataset the teacher is trained on. But current methods fail when the teacher parameters are more numerous than the training data, because the student overfits to the queries instead of aligning its parameters to the teacher. In this work, we explore augmentation techniques to best sample the input-output mapping of a teacher network, with the goal of eliciting a rich set of representations from the teacher hidden layers. We discover that standard augmentations such as rotation, flipping, and adding noise, bring little to no improvement to the identification problem. We design new data augmentation techniques tailored to better sample the representational space of the network's hidden layers. With our augmentations we extend the state-of-the-art range of recoverable network sizes. To test their scalability, we show that we can recover networks of up to 100 times more parameters than training data-points.

Problem

• Reverse-engineering neural network weights from limited input-output queries
• Overcoming student overfitting when teacher parameters exceed training data
• Developing tailored augmentations to sample hidden layer representations effectively

Innovation

• Data augmentation tailored to hidden layer representations
• Extending recoverable network size beyond training data limits
• Enabling weight reverse-engineering via optimized query sampling