Architecture-Aware Training Insights

• Architecture-aware training is a method where the model's structural features, such as computation graphs and routing topologies, explicitly influence optimization and hyperparameter settings.
• It leverages task-conditioned predictors, dynamic scheduling, and hardware-centric adjustments to improve efficiency, robustness, and deployment feasibility.
• The approach challenges one-size-fits-all recipes by advocating architecture-specific initialization and learning rate scaling, leading to more fair and accurate model evaluations.

Architecture-aware training denotes a family of learning procedures in which architectural structure is treated as an explicit conditioning variable of optimization rather than as a fixed backdrop for parameter fitting. In the literature, this conditioning may be imposed by the neural computation graph, task meta-features, hardware routing topology, quantization regime, student–teacher pairing, uncertainty objective, or robustness requirement. The resulting methods modify not only architecture search, but also initialization, learning-rate scaling, batch scheduling, sparsity control, latency-aware evaluation, and even the meaning of deployment feasibility itself (Kokiopoulou et al., 2019, Chen et al., 2024, Weber et al., 2024).

1. Conceptual scope and representations

A unifying feature of architecture-aware training is that architecture enters the optimization problem in a structured representation rather than as an opaque model identifier. In task-conditioned NAS, the architecture may be a continuous encoding $u:=\{\{\alpha\},\{\beta\},\{\gamma\}\}$, where softmax-relaxed operator choices $\alpha$, layer-gating variables $\beta$, and embedding-module selectors $\gamma$ make the architecture differentiable (Kokiopoulou et al., 2019). In learning-curve extrapolation, the architecture is a graph $G$ with adjacency matrix $A$ and node feature matrix $X$, encoded by message passing and pooled into a graph-level embedding that conditions latent ODE dynamics (Ding et al., 2024). In architecture-aware hyperparameter theory, the same object is formalized as a directed acyclic graph whose in-degrees and path counts determine admissible initialization and maximal stable learning rates (Chen et al., 2024).

The term also extends beyond the neural network graph in the narrow sense. In routing-aware neuromorphic training, the relevant architecture is the physical arrangement of Neuron Tiles and Routing Tiles, with hop distance $H(\text{NT}_i,\text{NT}_j)$ determining routing cost and memory demand (Weber et al., 2024). In accelerator co-design, architecture includes the compute pipeline, sparsity representation, and memory hierarchy; SPRING, for example, couples sparse computation, reduced precision, and monolithic 3D NVRAM to support both CNN training and inference (Yu et al., 2019). This suggests that “architecture-aware training” is not a single algorithmic template but a class of methods in which the optimization procedure is conditioned on structural constraints that materially affect learning dynamics or deployability.

A second common feature is the rejection of uniform recipes. Several papers argue directly that a one-size-fits-all schedule is inadequate. Architecture-aware scaling derives different maximal learning rates for different graph topologies and kernel sizes (Chen et al., 2024), DEBA calibrates batch-size adaptation thresholds per architecture (Belias et al., 5 Nov 2025), and subnet-aware supernet training assigns different decay schedules and momentum buffers to different subnet classes (Jeon et al., 13 Mar 2025). The architecture-aware perspective is therefore methodological as much as representational: it asserts that training rules should inherit structure from the model family they optimize.

2. Task-conditioned prediction and amortized architecture inference

A prominent line of work treats architecture-aware training as the learning of task-conditioned surrogates that predict architecture quality before expensive child-model training. “Fast Task-Aware Architecture Inference” builds an offline database of architecture–task–performance tuples $T=\{(z_i,u_i,v_i^*)\}_{i=0}^{M-1}$ and trains a deep value network $v(u,z;w)\approx v^*$, where $\alpha$0 is either a hand-crafted meta-feature vector or a learned permutation-invariant task embedding of the form $\alpha$1. Because the child architecture is continuously parameterized through softmax mixtures $\alpha$2, the online phase can optimize $\alpha$3 on a new task by simple gradient ascent, $\alpha$4, without training child models on the target task during inference (Kokiopoulou et al., 2019).

“Ranking architectures using meta-learning” retains the same task-tower/prediction-tower structure but replaces regression with a task-conditioned ranking objective. For architectures $\alpha$5 and $\alpha$6 on the same task, the linear ranking loss

$\alpha$7

and its quadratic variant train the network to preserve pairwise order rather than absolute validation accuracy. This aligns the surrogate with the downstream decision problem—selection of the best architecture—and the paper reports improved Spearman rank correlation and better architecture search results than $\alpha$8-based prediction (Dubatovka et al., 2019).

Few-label predictor training motivated a second development: architecture representation learning from unlabeled architectures. CAP encodes architectures as graphs, uses a GIN-style encoder for node updates,

$\alpha$9

and pre-trains the encoder with a context-aware self-supervised objective that matches a $\beta$0-hop central subgraph to its surrounding context. The downstream performance predictor is then fine-tuned with Bayesian Personalized Ranking,

$\beta$1

which the paper finds superior to MSE for ranking. CAP reports precise ranking with only 172 annotated architectures in NAS-Bench-101 and uses architecture structure itself as a pre-training signal (Ji et al., 2024).

Computation-aware encoding pursues a related objective at the representation level. CATE uses a Transformer encoder with architecture-dependent attention masks derived from adjacency or transitive-closure structure, and pre-trains on pairs of computationally similar architectures selected by a locality condition $\beta$2. Its pairwise masked-token objective learns embeddings that map architectures with similar accuracies to nearby regions in latent space, improving predictor-based search and generalization beyond the search space on which it was trained (Yan et al., 2021).

Architecture-aware prediction has also been extended from final accuracy to full training dynamics. LC-GODE combines a variational sequence encoder over early learning-curve prefixes with an architecture graph encoder and a graph-conditioned latent ODE,

$\beta$3

to extrapolate future epochs continuously while modeling uncertainty through a latent Gaussian posterior. The paper reports state-of-the-art extrapolation accuracy on both MLP and CNN learning curves, improved model-ranking quality, substantial regret reduction, and about a $\beta$4 speedup over exhaustive SGD training on the MLP tasks (Ding et al., 2024).

3. Architecture-conditioned optimization, schedules, and joint parameterization

Architecture-aware training is not restricted to predictor-based NAS. A distinct strand modifies the optimizer or its hyperparameters so that the first training step, or an online schedule, respects architectural geometry. “Principled Architecture-aware Scaling of Hyperparameters” derives topology-dependent initialization and maximal learning-rate laws from the criterion that each parameter should be maximally updated with the same mean squared change in pre-activations. For DAG networks, hidden-layer weights are initialized as

$\beta$5

while the output layer uses $\beta$6 scaling. The resulting maximal-update learning rate is

$\beta$7

for ReLU DAGs and

$\beta$8

for CNNs with kernel size $\beta$9. A major implication of the paper is evaluative: architecture-aware initializations and learning rates can substantially change network rankings in NAS benchmarks, so fixed training recipes may confound architectural comparison (Chen et al., 2024).

“Learning by Turning: Neural Architecture Aware Optimisation” moves the conditioning into the optimizer geometry itself. Nero constrains each neuron to be balanced,

$\gamma$0

and performs projected gradient descent with neuron-specific relative updates. Because the feasible set is a product of sphere-like constraints, the step size acquires an angular interpretation: each update rotates a neuron’s hyperplane rather than merely translating a weight vector. Nero is reported to train reliably without momentum or weight decay, to require little or no learning-rate tuning, and to have a memory footprint approximately square root that of Adam or LAMB (Liu et al., 2021).

Online schedule adaptation has received a similarly architecture-aware treatment. DEBA monitors gradient variance, gradient norm variation, and loss variation,

$\gamma$1

and adjusts the batch size by increase, rollback, or hold decisions using architecture-specific thresholds. Its empirical conclusion is explicitly anti-universalist: lightweight and medium-depth architectures such as MobileNet-V3, DenseNet-121, and EfficientNet-B0 can realize $\gamma$2 training speedups with simultaneous accuracy gains, whereas ViT-B16 shows only minimal speedup and ResNet-50 exhibits high variance and occasional degradation (Belias et al., 5 Nov 2025).

A more radical formulation makes architecture and weights a single optimization object. “Simultaneous Weight and Architecture Optimization for Neural Networks” trains a multi-scale encoder–decoder so that functionally similar MLPs map to nearby latent points even if their depths and weights differ. A latent code $\gamma$3 is decoded into a concrete architecture–weight pair $\gamma$4, and $\gamma$5 is optimized directly on task loss plus sparsity regularization,

$\gamma$6

with $\gamma$7. Architecture selection and parameter fitting thus occur simultaneously in latent space rather than in a search-then-retrain pipeline (Huang et al., 2024).

4. Hardware, quantization, and deployment-driven co-design

In deployment-oriented work, architecture-aware training makes hardware feasibility part of the learning objective. The routing-aware neuromorphic case is explicit: a network that is accurate in software may be unusable on event-based hardware if its connectivity cannot be placed and routed efficiently. The routing-aware DeepR extension defines a hop-dependent sparsity profile

$\gamma$8

and constrains pruning and random reassignment so that $\gamma$9 and $G$0. On SHD with the Mosaic small-world routing architecture, the paper reports about 5 percentage points higher accuracy at fixed memory count and the same maximum accuracy of 68.5% with one order of magnitude fewer memory elements than a non-routing-aware baseline (Weber et al., 2024).

Quantization-aware architecture search pursues a similar co-optimization under extremely low precision. OQAT argues that the best floating-point architecture is often not the best quantized architecture and therefore trains a quantized supernet in a single stage. Uniform quantization is parameterized by a learned step size $G$1,

$G$2

with one shared step size per layer, and lower-bit supernets are initialized by bit inheritance,

$G$3

The paper reports that OQAT-2bit-M reaches 61.6%/61.7% ImageNet Top-1 accuracy, outperforming 2-bit MobileNetV3 by about 9% while using 10% less computation cost, and that total GPU hours for $G$4 deployed models drop to 1.2k versus 10.8k for SPOS and 11.4k for BMobi (Shen et al., 2020).

At the accelerator level, SPRING embodies architecture-aware training as a joint choice of sparse data representation, arithmetic format, and memory system. It encodes activations and weights with binary masks, filters dangling non-zeros before MAC execution, and employs stochastic rounding,

$G$5

so that reduced-precision training remains accurate. Coupled to a monolithic 3D RRAM interface, SPRING reports geometric-mean improvements of $G$6 performance, $G$7 power reduction, and $G$8 energy efficiency for CNN training relative to GTX 1080 Ti (Yu et al., 2019).

A neighboring literature treats architecture-aware training as learned design-space exploration for hardware itself. AIRCHITECT quantizes accelerator design and mapping choices into labels and trains an embedding-plus-MLP recommender that predicts the optimal label from workload dimensions and design constraints; across three systolic-array case studies it reports up to 94.3% validation/test accuracy and predicted configurations achieving 99.9% of the best possible performance on $G$9 GEMM workloads (Samajdar et al., 2021). PEL-NAS extends the hardware-aware viewpoint to LLM-driven NAS by partitioning the search space into complexity niches, co-evolving prompts and architectures, and scoring candidates with an XGBoost surrogate on 13 zero-cost proxies; on HW-NAS-Bench it reports higher HV, lower IGD, and up to 54% lower latency than baselines at similar accuracy, with search cost reduced to 3 minutes via API calls (Zhu et al., 1 Oct 2025).

5. Supernet-aware search, distillation, and uncertainty-aware objectives

Weight-sharing NAS exposes another architecture-aware failure mode: different subnets do not optimize at the same rate, and a single training policy biases the supernet. “Subnet-Aware Dynamic Supernet Training for Neural Architecture Search” identifies two problems—unfairness toward low-complexity subnets and noisy momentum—and addresses them with a complexity-aware learning-rate scheduler,

$A$0

and momentum separation over subnet clusters,

$A$1

Across MobileNet and NAS-Bench-201 spaces, the method improves both Kendall’s Tau and Top-1 accuracy over SPOS, FairNAS, and FSNAS, while remaining a drop-in modification for sampling-based NAS (Jeon et al., 13 Mar 2025).

RL-based NAS has been made architecture-aware by explicit knowledge reuse. FNAS augments PPO with an uncertainty-aware critic that predicts both value $A$2 and its own error magnitude through

$A$3

an architecture knowledge pool for block-level parameter reuse via similarity matching, and an architecture experience buffer with prioritized replay,

$A$4

The paper reports about $A$5 reduction in RL-based NAS cost on MNAS-style search, from about 20,000 GPU hours to about 2,000 GPU hours, while improving or matching downstream performance on ImageNet, face recognition, and COCO transfer (Liu et al., 2021).

Architecture-aware training has also been applied to knowledge distillation. TRADE searches student architectures under the KD loss

$A$6

using trust-region Bayesian optimization over a search space containing MBConv, Fused-MBConv, and Transformer blocks. The scoring function jointly considers accuracy, parameter count, FLOPS, and latency, and the paper reports that the resulting students are consistently smaller and faster than fixed baselines while achieving competitive or better KD accuracy; for example, TRADE-ViT reduces FLOPS from 4.2B to 1.8B and latency from 18.1 ms to 4.5 ms while reaching 81.5 top-1 under KD (Kim et al., 2022).

Uncertainty itself can be made an architectural search target. $A$7DARTS inserts concrete dropout into DARTS cells, augments the inner objective with a Monte-Carlo regularizer,

$A$8

and adds predictive variance to the validation objective. The method searches for architectures that are both accurate and low-uncertainty; on CIFAR-10 it reports 96.22% accuracy with uncertainty 0.162, compared with 94.89% and 0.396 for RDARTS-CD (Chakraborty et al., 2021).

6. Invariance, robustness, evaluation, and limitations

Architecture-aware training is also used to encode symmetries that training-only methods do not guarantee. “Training or Architecture? How to Incorporate Invariance in Neural Networks” proposes orbit mapping, in which a canonical representative is selected from the transformation orbit

$A$9

and then fed to an otherwise standard network. For equivariance, the architecture canonicalizes and then maps back,

$X$0

The paper argues that this architecture-based approach yields provable invariance for suitable group actions, can handle continuous transformations, and is often more computationally efficient than augmentation or adversarial training (Gandikota et al., 2021).

Robustness to adversarial attacks has likewise been reframed as an architectural property rather than solely a training property. RobArch conducts a fixed-parameter-budget study of more than 150 ImageNet models under adversarial training and distills 18 design guidelines. Among the reported conclusions are that 4-stage networks are best among the tested stage counts, the depth pattern $X$1 is optimal in the fixed-budget setting, inverted bottlenecks and large dilation hurt robustness, and SE plus SiLU materially help. The resulting RobArch-L reaches 48.94% AutoAttack accuracy on the RobustBench ImageNet leaderboard (Peng et al., 2023).

These results have direct consequences for evaluation methodology. Architecture-aware scaling shows that architecture-specific initialization and learning rates can alter benchmark rankings, so a fixed recipe may reward architectures that happen to match default hyperparameters rather than architectures that are intrinsically superior (Chen et al., 2024). LC-GODE reaches a parallel conclusion for early-stopping settings: learning-curve ranking improves when curve extrapolation is conditioned on architecture rather than treated as pure time-series prediction (Ding et al., 2024). A broader implication is that benchmark fairness depends on whether the training protocol equalizes optimization opportunity across heterogeneous architectures.

The limitations of architecture-aware training are recurrent and substantial. Task-conditioned predictors require many related source tasks and a historical database of architecture training results, and their continuous inference procedures presume a differentiable architecture parameterization (Kokiopoulou et al., 2019). Subnet-aware supernet methods are designed for sampling-based NAS and are not directly applicable to differentiable NAS (Jeon et al., 13 Mar 2025). DEBA depends on architecture-specific threshold calibration and derives its gains primarily for moderately stable architectures rather than universally (Belias et al., 5 Nov 2025). Hardware-aware methods often rely on case-study assumptions, such as a proxy mappability function, a particular hardware lookup table, or a fixed routing model (Weber et al., 2024). The literature therefore supports a precise but narrow conclusion: architecture-aware training is most effective when the structural variables that govern optimization, generalization, or deployment can be represented explicitly enough to enter the learning algorithm itself.