AutoML-Agent Framework

• AutoML-Agent Framework is a modular, decentralized multi-agent system that automates ML pipeline development including data augmentation, NAS, and HPO.
• It employs agents interacting via structured protocols such as MDPs and MARL, achieving high validation accuracy through off-policy actor-critic methods.
• The framework balances accuracy with computational cost using reward shaping and counterfactual credit assignment to optimize convergence in high-dimensional search spaces.

An AutoML-Agent Framework, sometimes referred to as a multi-agent AutoML system, is a modular, decentralized software architecture in which specialized, interacting agents jointly automate the design, optimization, and validation of machine learning pipelines. In contemporary research, such frameworks are realized by representing key machine learning modules (e.g., data augmentation, neural architecture search, hyperparameter optimization) as independent agents that cooperate through structured protocols—often formulated as Markov decision processes (MDPs) and solved using multi-agent reinforcement learning (MARL) or policy gradient methods. The multi-agent paradigm fundamentally improves credit assignment, joint search efficiency, and convergence in complex AutoML tasks, particularly in high-dimensional search spaces and under computational constraints (Wang et al., 2022).

1. Core Architecture and Agent Decomposition

A canonical AutoML-Agent Framework, as exemplified by MA2ML (Wang et al., 2022), decomposes the full AutoML pipeline into distinct modules, each controlled by a dedicated agent. In MA2ML, three agents correspond to:

• AUG Agent: Selects a data augmentation policy from a discrete library (e.g., 25 sub-policies, each parameterized by operation, magnitude, and probability).
• NAS Agent: Chooses neural network architectural hyperparameters within large spaces (e.g., kernel sizes, channel widths, depths, input resolution) using NASNet or FBNetV3-like search domains.
• HPO Agent: Samples optimizer type, learning rate schedule, weight decay, and other hyperparameters (mixup, dropout, stochastic depth, EMA).

Each agent's policy is implemented as a recurrent controller, typically a single-layer LSTM with 100 hidden units, which defines a mapping from a (possibly trivial) internal state to an action in the agent's action space. The actions are chosen jointly at each iteration, forming a one-step MDP with an immediate reward computed as the top-1 validation accuracy after training the selected pipeline. The key joint objective is

$$J(\Theta) = \mathbb{E}_{A_1,\ldots,A_n \sim \pi}[R],$$

where $\Theta = \{\theta_i\}_i$ denotes all agent parameters and $R$ is the scalar reward.

2. Reward Design and Budget-Aware Objectives

The system reward in the AutoML-Agent framework is generally tied to the validation performance of the pipeline assembled by the agents. In resource-constrained scenarios, explicit computational cost components (e.g., FLOPs, model size) are incorporated. MA2ML, for ImageNet tasks, uses

$$R(m) = \mathrm{Acc}(m) \cdot \left[\frac{\mathrm{FLOPs}(m)}{F_0}\right]^w, \quad w = -0.07,$$

where $F_0$ is the target FLOPs budget, biasing the agents toward computationally efficient solutions. This transforms the reward landscape to account for both accuracy and resource requirements.

3. Marginal Contribution and Credit Assignment

One of the central algorithmic challenges in multi-agent AutoML—absent in monolithic frameworks—is attributing the joint reward back to the individual agents' actions. MA2ML addresses this using a centralized critic $Q(S, A_1, ..., A_n)$ and a counterfactual baseline to compute each agent's marginal credit:

• Counterfactual Baseline for agent $i$:

$$b(S, A_{-i}) = \mathbb{E}_{A_i \sim \pi_i}[Q(S, A_i, A_{-i})]$$

• Advantage for the agent's actual action:

$\Delta R_i = Q(S, A_1, ..., A_n) - b(S, A_{-i})$

This mechanism avoids the confounding typical in global reward signals and provides shaped, agent-specific feedback suitable for efficient actor-critic updates.

4. Off-Policy Actor-Critic and Optimization Guarantees

Sample efficiency is improved via off-policy actor-critic learning, employing experience replay and enforcing a Kullback–Leibler divergence regularization between the “behavior” and “target” policies for each agent. The global regularized objective is

$$J(\pi, \rho) = \mathbb{E}_{A \sim \pi}[Q(S, A)] - \lambda D_{\mathrm{KL}}(\pi(\cdot|S) \,\|\, \rho(\cdot|S))$$

The training alternates between collecting new joint actions and rewards (pipeline evaluations), critic updates (fitting $Q$ to observed rewards), and actor updates (policy gradients using the credit assignment). MA2ML proves monotonic improvement of the unregularized joint objective under a divergence-regularized policy iteration scheme. Specifically:

• If at iteration $k$

$$\pi^{k+1} = \arg\max_{\pi} \mathbb{E}_{A \sim \pi}[R - \lambda \log(\pi / \rho^k)]$$

and

$\rho^{k+1} = \pi^k,$

then $J(\pi^{k+1}) \ge J(\pi^k)$ and the sequence converges to a (local) optimum.

5. Large-Scale Empirical Results

MA2ML is benchmarked on CIFAR-10/100 and ImageNet with varying FLOPs ceilings. The search space is decomposed as AUG ($\sim\!10^{12}$), NAS ($\sim\!10^{37}$), and HPO ($\sim\!10^{12}$). Notable results include:

Dataset	Method	Top-1 (%)	FLOPs (M)
CIFAR-10	MA2ML	97.77	—
CIFAR-10	MA2ML-Lite	97.70	—
CIFAR-100	MA2ML	85.08	—
CIFAR-100	MA2ML-Lite	84.80	—
ImageNet	MA2ML-A	79.3	490
ImageNet	MA2ML-B	79.7	596
ImageNet	MA2ML-C	80.1	694

MA2ML-B achieves 79.7% top-1 accuracy under 600M FLOPs, outperforming prior RL and gradient-based NAS and joint AUG+NAS methods (DAAS, DHA, FBNetV3). The benefit of the full credit assignment and off-policy learning, relative to on-policy or REINFORCE ablations (MA2ML-Lite), amounts to an additional +0.6% at fixed computational budgets.

6. Relationship to Other Multi-Agent and MARL-Based AutoML Approaches

By formalizing joint AutoML as a multi-agent MDP, the AutoML-Agent paradigm generalizes single-agent RL (NAS, HPO) and enables modular, scalable joint search over combinatorially large pipeline spaces. The explicit agent design allows for targeted improvements in sample efficiency, reward shaping, and credit assignment—all of which are critical at the scale of contemporary AutoML search spaces. The use of divergence-regularized off-policy actor-critic algorithms and counterfactual advantage estimation are distinct contributions relative to earlier monolithic or hand-crafted AutoML controllers (Wang et al., 2022).

---

References:

• “Multi-Agent Automated Machine Learning” (Wang et al., 2022)