Beyond Model Base Retrieval: Weaving Knowledge to Master Fine-grained Neural Network Design

TL;DR

The paper tackles the inefficiency of neural architecture search and the suboptimality of model retrieval by introducing M-DESIGN, a retrieval-augmented refinement framework that uses fine-grained Architecture Modification-Gain Graphs to encode how small architectural edits impact performance. It weaves evidence across related tasks through a Bayesian dynamic task similarity model and online updates, enabling adaptive retrieval of edits and predictive planning for multi-hop gains under limited budgets. Predictive task planners fill gaps when evidence is missing or when distribution shifts occur, allowing extrapolation beyond observed edits. Across 33 task–data scenarios on graph data with 67,760 trained GNNs over 22 datasets, M-DESIGN achieves the best design-space performance in 26 cases, with low computational overhead, demonstrating improved efficiency and robustness for fine-grained neural network design.

Abstract

Designing high-performance neural networks for new tasks requires balancing optimization quality with search efficiency. Current methods fail to achieve this balance: neural architectural search is computationally expensive, while model retrieval often yields suboptimal static checkpoints. To resolve this dilemma, we model the performance gains induced by fine-grained architectural modifications as edit-effect evidence and build evidence graphs from prior tasks. By constructing a retrieval-augmented model refinement framework, our proposed M-DESIGN dynamically weaves historical evidence to discover near-optimal modification paths. M-DESIGN features an adaptive retrieval mechanism that quickly calibrates the evolving transferability of edit-effect evidence from different sources. To handle out-of-distribution shifts, we introduce predictive task planners that extrapolate gains from multi-hop evidence, thereby reducing reliance on an exhaustive repository. Based on our model knowledge base of 67,760 graph neural networks across 22 datasets, extensive experiments demonstrate that M-DESIGN consistently outperforms baselines, achieving the search-space best performance in 26 out of 33 cases under a strict budget.

Figures (9)

• Figure 1: Illustration on the effectiveness-efficiency dilemma of conventional neural architecture search and model retrieval methods (Upper part), and the novel framework of M-DESIGN that adaptively weaves fine-grained modification-gain information for a better trade-off (Lower part).
• Figure 2: Illustration of transferability estimators in model retrieval: 1) ranking-based, 2) embedding-based, and 3) semantic-based.
• Figure 3: Overview of M-DESIGN. The knowledge base stores graph-structured modification evidence (white). A knowledge weaving engine (yellow) aggregates gains from related tasks and updates a Bayesian task-similarity belief online using observation.
• Figure 4: Long-run model refinement trajectories of M-DESIGN compared to search-based and retrieval-based methods on Cornell.
• Figure 5: Bar charts showing the empirical supports for the theoretical assumption on Cora and Cornell. Each benchmark dataset has three bars: our stabilized task similarity snapshot, R-squared (Linearity), and Shapiro-Wilk p-value (Gaussian).

Theorems & Definitions (3)

• Definition 3.1: Adaptive Model Retrieval Problem
• Definition 4.1: Task Similarity via Local Modification Consistency
• Definition 4.2: Knowledge Weaving for Adaptive Retrieval