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

1. Problem Statement

Designing high-performance neural networks for new tasks requires balancing optimization quality with search efficiency. Current methodologies face a fundamental trade-off:

• Neural Architecture Search (NAS): While capable of finding near-optimal structures, traditional NAS relies on exhaustive trial-and-error, making it computationally expensive and inefficient ("reinventing the wheel").
• Model Retrieval: Recent approaches retrieve pre-trained models from extensive repositories to bypass searching from scratch. However, these methods often yield suboptimal static checkpoints. They struggle to adapt to the fine-grained architectural edits (e.g., changing message-passing mechanisms in Graph Neural Networks) required to match idiosyncratic data properties and distribution shifts.

Existing retrieval methods typically compute task similarity once and keep it static. This fails during sequential refinement, especially under Out-of-Distribution (OOD) shifts, where the transferability of prior evidence evolves as the architecture changes. Furthermore, when direct evidence is missing or the new task is significantly OOD, simple retrieval becomes unreliable or misleading.

2. Methodology: M-DESIGN

The authors propose M-DESIGN, a retrieval-augmented framework that moves beyond static model retrieval to adaptive knowledge weaving. The core idea is to model performance gains induced by fine-grained architectural modifications as "edit-effect evidence" and weave this evidence dynamically to discover near-optimal modification paths.

2.1 Architecture Modification-Gain Graph

Instead of storing coarse performance records of complete models, M-DESIGN constructs a Model Knowledge Base (MKB) organized as an Architecture Modification-Gain Graph ($G_\Delta$).

• Nodes: Represent architecture configurations.
• Edges: Represent feasible 1-hop modifications.
• Weights: Represent the historical performance gain ($\Delta P$) of a specific modification on a specific benchmark task.
  This structure allows the system to reason about modification trajectories and chain local edit patterns across different tasks.

2.2 Adaptive Knowledge Weaving

The refinement process is framed as finding a sequence of modifications that maximizes expected gain based on historical evidence.

• Fine-grained Knowledge Weaving: The system approximates the gain on an unseen task ($D_u$) by aggregating fine-grained gain evidence from benchmark tasks ($D_b$). It assumes that tasks with consistent local modification behavior share similar refinement landscapes.
• Dynamic Task Similarity (Bayesian Update): A critical innovation is the adaptive retrieval mechanism. Instead of using a static similarity metric, M-DESIGN maintains a Bayesian belief over task similarity ($S_t(D_u, D_i)$).
  • As the model is refined and new gains are observed on the unseen task, the system updates the similarity belief online.
  • If new evidence contradicts the initial assumption that a benchmark task is similar, the posterior belief downweights that task, preventing overconfidence and correcting for distribution shifts.
  • This is modeled using a Gaussian likelihood where the gain on the unseen task is linearly related to the gain on the benchmark task, scaled by a transfer factor.

2.3 OOD Adaptation with Predictive Task Planners

To handle scenarios where direct evidence retrieval is unreliable (due to sparse historical coverage or significant OOD shifts), M-DESIGN employs Predictive Task Planners.

• These are regression models (instantiated as edge-regression GNNs) trained on the modification-gain graphs of benchmark tasks.
• They extrapolate gains for modifications that lack direct historical records (missing edges) or when the similarity belief drops below a threshold.
• The system uses a replay buffer of observed interactions to fine-tune these planners online, ensuring synthetic evidence aligns with the target distribution over time.

3. Key Contributions

1. Reframing Model Refinement: The authors propose shifting from static model retrieval to adaptive retrieval, which discovers near-optimal modification paths using repository priors while dynamically calibrating transfer evidence to reduce reliance on repository completeness.
2. Bayesian Dynamic Task Similarity: They introduce an iterative retrieval-and-refine algorithm that updates transfer beliefs online using observed gains. This includes training predictive task planners to extrapolate reliable evidence under OOD shifts.
3. Comprehensive Empirical Validation: The work demonstrates strong performance across 33 task-data pairs (covering node classification, link prediction, and graph classification). Under a strict budget of 100 model evaluations, M-DESIGN achieves the search-space best performance in 26 out of 33 cases.
4. Open Model Knowledge Base: The authors release a structured MKB containing 67,760 trained Graph Neural Networks across 22 datasets and 3 task types, explicitly recording edit-effect records.

4. Experimental Results

• Effectiveness: M-DESIGN consistently outperforms strong baselines, including search-based methods (Random, RL, EA, GraphNAS) and retrieval-based methods (GraphGym, NAS-Bench-Graph, AutoTransfer, DesiGNN). It reaches the global optimum in the design space for the majority of tested pairs.
• Efficiency: The method adds negligible overhead (< 0.44 seconds per iteration) compared to model evaluation time. It reaches high-performance levels significantly faster than baselines, many of which fail to reach the target performance even after 100 evaluations.
• Mechanistic Validation:
  • Linearity & Normality: Empirical analysis confirms that modification gains transfer approximately linearly between similar tasks and follow a Gaussian distribution, validating the Bayesian update assumptions.
  • Dynamic vs. Static: Ablation studies show that removing the dynamic similarity update causes a significant drop in performance and alignment with ground-truth local consistency.
  • OOD Adaptation: The predictive task planners are crucial for OOD-sensitive datasets (e.g., Cornell, Texas), improving refinement stability where static transfer fails.
• Generalization: The framework was also tested on tabular (HPOBench) and image (NATS-Bench) domains, showing that while refinement is highly valuable for data-sensitive tabular tasks, one-shot retrieval is often sufficient for relatively homogeneous image tasks.

5. Significance and Claims

The paper claims that M-DESIGN resolves the dilemma between the high cost of exhaustive search and the suboptimality of static retrieval. By treating model refinement as a continuous process of weaving fine-grained modification evidence, the method achieves a better trade-off between efficiency and effectiveness.

The authors emphasize that their approach is particularly significant because it addresses the evolving transferability of knowledge. Unlike previous methods that assume static task similarity, M-DESIGN acknowledges that as an architecture evolves, the relevance of different source tasks changes. The integration of predictive planners further allows the system to function robustly even when the repository does not cover the specific modifications needed for a new task, effectively bridging the gap between in-distribution retrieval and out-of-distribution generalization.

The work concludes that while the method relies on the validity of local gain modeling and repository coverage, it provides a systematic, sample-efficient alternative to ad-hoc trial-and-error, reducing the energy costs of exhaustive search. The authors note limitations regarding extreme OOD tasks or total absence of evidence, suggesting future work in scaling to larger model families and uncertainty-aware gain prediction.