Combining Multi-Objective Bayesian Optimization with Reinforcement Learning for TinyML

In plain words

Microcontrollers have very little memory and compute, so a neural network for one has to be small and fast. Finding a good one means choosing many settings at once: how long to train, the learning rate, how much of the network to prune. Each trial is a full training run, so the search has to find good trade-offs in few trials.

The paper's approach is Bayesian optimization. It keeps a model of the trials so far and uses that model to pick the next trial. Here, the next trial is chosen by a small group of simple reinforcement-learning agents. Each agent searches the model's landscape from a different good starting point. The agents use a simple random-search method rather than a heavy learning algorithm. The output is a set of networks that trade accuracy against memory and compute, and each one can be deployed without retraining.

TL;DR

• Targets: DNNs for microcontrollers with fixed flash, RAM and FLOPs limits, tuned for accuracy and resources at the same time §1.
• Method: multi-objective Bayesian optimization with a Gaussian-process surrogate. The acquisition step is replaced by an ensemble of ARS agents trained on GP posterior samples §4 alg. 1.
• Search: training, pruning and quantization hyperparameters, with a budget of 150 samples §5 tab. 1.
• Result: author-reported higher hypervolume than ParEGO, TurBO, MorBO, NSGA-II and random search on two use cases §5.1 fig. 2. Code is an MIT-licensed Zenodo archive, not read for this card.

Problem

Deploying DNNs on microcontrollers is limited by the target's flash and RAM. The introduction gives the limits as under 1 MB of flash and under 512 Kb of RAM, and the experiments use 256 KB of RAM, 1 MB of ROM and 1e9 FLOPs §1 §5.

The paper's view of earlier work is that black-box NAS methods are reliable but slow, differentiable NAS is prone to overfitting the validation set, and zero-cost NAS gives only proxy statistics such as FLOPs §1. Multi-objective Bayesian optimization is sample-efficient, but the paper argues it has not been used much for microcontroller NAS §2. Each trial trains a network, so the search budget is tight: 150 samples per run §5.1.

Multi-objective goal: find an approximate feasible Pareto front within a budget of n evaluations. Quality is measured by the hypervolume indicator with a reference point taken from domain knowledge §3.

Method

Background. Expected improvement is the usual acquisition function. It is written as an expectation over the posterior and approximated with Monte-Carlo samples §3 eq. 1 §3 eq. 2:

EI(x) = \mathbb{E}\max(f(x) - f^*, 0) §3 eq. 1

qEI(x) \approx \frac{1}{N}\sum_{i=1}^{N}\max_{j}\{\max(\xi_i - f^*, 0)\}, \quad \xi_i \sim P(f(x) \mid \mathcal{D}) §3 eq. 2

Overview. The outer loop fits GP surrogates on all evaluated trials. The proposed architectures are said to be deployable directly, without further training §1. For each sample, it trains a set of ARS agents, each starting from a different point on the feasible Pareto front. The agents' action space is the search hyperparameters. Each agent is trained on rewards computed from posterior samples of the GP, so no real DNN training happens inside the RL loop §4 alg. 1.

Numbered steps:

1. Initial prior: propose the first 10 candidates with Latin-hypercube sampling and evaluate them §4 §5.1.
2. Fit a GP to each objective by maximising the marginal likelihood of the evaluations so far alg. 1 §4.
3. Choose start states: if the feasible Pareto front has more points than there are agents, run k-means on it to pick diverse starts §4 alg. 1. The number of agents is not stated.
4. Draw a random scalarization weight vector λ for this sample §4.
5. For each agent, train an ARS policy. Rewards come from quasi-Monte-Carlo samples of the GP, using fixed base samples §4. The reward is defined in eq. 7 and eq. 8.
6. Roll out each trained policy from its start state and keep the best rollout §4 alg. 1.
7. Evaluate the proposed configuration: train the DNN, prune it, quantize it, and generate code from it to measure memory §4. Add the result to the prior.

Scalarization and reward. The n objectives are reduced to one value with an augmented Chebyshev scalarization. The reward is the change in that value relative to the best value seen so far, with negative changes scaled down §4:

f_\lambda(x) = \max_{j=1}^{n}(\lambda_j f_j(x)) + \rho \sum_{j=1}^{n} \lambda_j f_j(x), \quad \rho = 0.005 §4 eq. 7

r(x) = \begin{cases} f_\lambda(x) - f_\lambda(f^*) & \text{if } f_\lambda(x) - f_\lambda(f^*) > 0 \\ 10^{-3}\,(f_\lambda(x) - f_\lambda(f^*)) & \text{otherwise} \end{cases} §4 eq. 8

The paper's argument is that EI clamps negative improvements to zero, which favours uncertain regions. Scaling negative improvements by a small factor keeps the trade-off while letting the policy learn from worse points §4.

ARS inner loop alg. 1 §3. ARS is a derivative-free random-search method for learning policies. The policy is parameterized by θ. Each iteration samples N random directions with i.i.d. standard normal entries, evaluates each direction in both signs over horizon H, keeps the best b directions, and takes a step along their reward differences. The paper's text sets ν = 0.008 and reports a top-k rate of 1% §3 §5.2.

# Multi-objective BO with ARS agents  [alg. 1]
# Input:  prior (evaluated trials), budget J, objective f
# Output: prior extended with new trials

def ars_mobopt(prior, budget, objective, agents):
    for iteration in range(budget):
        # Fit one GP per objective on the trials so far  [alg. 1] [§4]
        surrogates = fit_gaussian_processes(prior)
        # Start states: k-means centres of the feasible front  [§4] [alg. 1]
        starts = kmeans_centres(feasible_front(prior), agents)
        # Scalarization weights for this sample, drawn at random  [§4]
        weights = draw_random_weights()
        trained = []
        for start in starts:
            # Train one ARS agent on rewards from GP posterior samples  [§4]
            trained.append(train_ars_agent(start, surrogates, weights))
        # Roll out each trained agent and keep its best rollout  [alg. 1] [§4]
        proposals = [best_rollout(agent) for agent in trained]
        # Evaluate proposals on the real objective and extend the prior  [§4]
        for proposal in proposals:
            prior.append(objective(proposal))
    return prior


def train_ars_agent(start, surrogates, weights):
    # Linear policy in the paper's algorithm; text uses 2 layers of 64 [§4]
    theta = initial_policy(start)  # initialisation not stated
    while not ars_stop_condition():  # stopping rule not stated  [alg. 1]
        # Sample N directions with i.i.d. standard normal entries  [alg. 1]
        directions = random_normal(N, *theta.shape)
        # Roll out each direction in both signs over horizon H  [alg. 1] [§4]
        plus, minus = rollouts(theta, directions, nu, H, surrogates, weights)
        # Keep the b directions with the largest max(plus, minus)  [alg. 1]
        top = best_directions(directions, plus, minus, b)
        # Step along the reward-weighted sum of the kept directions  [alg. 1]
        theta = theta + alpha / (b * reward_std(top)) * update_sum(top)
    return theta

Note on algorithm symbols: the algorithm box uses v for the exploration noise and N for the number of directions. The text uses ν and N. The algorithm also uses p × n direction matrices alg. 1, which conflict with the two-layer policy described in the text §4.

Hyperparameters

Search space for the DNN training and pruning hyperparameters, the action space of the ARS agents tab. 1:

Fixed settings of the search and the models:

Results

All results are author-reported. Hardware used to train the networks is not stated in the text.

Algorithmic benchmark §5.1 fig. 2. Two tasks, 150 samples, 5 seeds. Methods: ARS-MOBOpt against ParEGO, TurBO, MorBO, NSGA-II and random sampling. The text reports that ARS-MOBOpt had the highest hypervolume after about 50 to 75 samples on both tasks, and that all Bayesian methods beat NSGA-II. It gives no numeric hypervolumes; the values are only in the figure.

Time-series benchmark §5.1 tab. 4. ARS against ParEGO on 21 multivariate UEA & UCR datasets. Table 2 is the excerpt of the datasets where accuracy reached 70%; the full table is below. Bold marks in the source are omitted.

In the source, the resource columns of the last 11 rows (from HandMovementDirection down) are grey. They report the highest-accuracy network, because no network reached 70% tab. 4 tab. 2. The text claims ARS is better than ParEGO on 7 of 10 datasets with at least 70% accuracy, and on par on one §5.1. Table 4 does not support that count for HV, see Pitfalls.

MCUNet and zero-shot comparison on CIFAR-10 §5.4.2 tab. 3. The table lists five MCUNet models taken from the MCUNet model list, one ProxylessNAS model searched for CIFAR-10, and three ARS-MOBOpt designs. Parentheses in the source give the memory under the paper's own mapping tool. MACs in the table are as printed §5.4.2.

Other reported numbers in this part: ProxylessNAS estimated 86.53% and 27.16M FLOPs for its model, and the retrained model reached 82.33% at 25.5M §5.4.2. The paper reports that its designs use less ROM than MCUNet, at higher compute, and that ours-large is slightly more accurate than mcunet-in4 §5.4.2.

PPO comparison on ResNet-18/CIFAR-10 and MobileNetV3/DaLiAc §5.4.1 fig. 5. With default PPO settings, the Hypervolume was no better than random sampling on ResNet. After tuning with Optuna (20 PPO configurations, each run on 20 samples with a fixed prior of 2), the tuned PPO beat ARS on ResNet §5.4.1 fig. 5. The same tuned settings barely beat random sampling on MobileNetV3/DaLiAc §5.4.1. Authors' conclusion: PPO needs per-problem tuning, and ARS was chosen for simplicity §5.4.1. The paper's values for the tuned and default PPO settings are in app. C.

AMC comparison on MobileNet/CIFAR-10 §5.4.3 fig. 6. Single-run comparison: AMC gives one trade-off, while ARS-MOBOpt gives a front after 50 samples. The paper reports that AMC matches the best accuracy, but ARS-MOBOpt has clearly better RAM and FLOPs trade-offs §5.4.3. No numbers are in the text.

Ablation of ARS hyperparameters §5.2 fig. 3. On MobileNetV3/DaLiAc, the paper reports that about 3000 directions, a top-k of 1%, and H = 4 gave the best results. Few directions performed poorly. Longer horizons had higher variance §5.2.

Paper vs code

The paper states the code is an open-source Python implementation of the solver, including the ARS and PPO agents. The link is in the footnote of the Introduction: Zenodo record 10.5281/zenodo.13918411 §1. That record was checked on Zenodo: authors Mark Deutel, Georgios Kontes, Christopher Mutschler, Jürgen Teich; MIT license; version 1, dated October 2024; one file, paper_release.zip, which contains the paths deepcompression/rl/ars/ (environment.py, sampler.py), deepcompression/apps/run_optimization.py, deepcompression/optimizer.py, deepcompression/config.py, scripts/optim_*.sh and training/config/*/optim.yml. No GitHub repository was found by search. The code is not on GitHub, so it was not cloned.

Not checked: the archive could not be extracted during this compile (extraction was not permitted), so the code was not compared with the paper line by line. The values above come from the paper alone. Before this card is marked reviewed, someone must compare at least the ARS update, the reward, the search space and the number of agents with the archive. Until then, treat the details above as unverified against the code. code_commit is left empty because the Zenodo record has no commit; lint warns about this.

Pitfalls

• Table 2 vs text. The text claims ARS beats ParEGO on 7 of 10 datasets, with one on par §5.1. Using the HV column of the table, ARS is higher on 6 datasets, equal on 2 (UWaveGestureLibrary, Heartbeat), and lower on 2 (BasicMotions, PEMS-SF) tab. 4. Using accuracy, it is higher on 7, equal on 2 and lower on 1. The text does not say which metric it means.
• Units of memory. The table headers say ROM [Kb] and RAM [Kb], but values such as 120512 are only consistent with the 1 MB and 256 KB limits if they are bytes tab. 2 tab. 4 §5. Treat the unit as unconfirmed.
• RAM limit in the MCUNet table. ours-large uses 276 KB of RAM, more than the 256 KB RAM constraint stated in §5. The text's claim that the paper's networks beat all MCUNet variants on accuracy holds only for ours-large against mcunet-in4 (87.96 vs 87.60) tab. 3 §5.4.2.
• Rollout horizon. The text says H = 4 is best §5.2. The alt text (\Description) of fig. 3 says 16 fig. 3. The caption of fig. 3 gives the model as MobileNetv3 with 1.6M initial parameters, while §5 gives 2.8M for MobileNetV3 and 1.6M for ResNet-18 fig. 3 §5.
• Objective direction. §3 is written for maximization. Accuracy should be maximized, but ROM, RAM and FLOPs should be minimized. The paper does not say how those are sign-flipped or normalised before eq. 7 §3 §5.1. The max term in eq. 7 is the one written in the paper; check it against the code before implementing.
• Proposals per iteration. §4 says one policy, the one with the best reward, proposes the next trial. Algorithm 1 evaluates the best rollout of every agent each iteration. The number of real evaluations per iteration therefore differs (1 vs L). The text does not say which the experiments used §4 alg. 1.
• Policy form. The algorithm uses linear policies with direction matrices in R^(p×n) alg. 1. The text uses two linear layers of width 64 §4. Implement the text form unless the code says otherwise.
• ARS update details. The update uses σ_R, "the standard deviation of the rewards used in the update step" alg. 1. Which rewards those are is not stated (the kept top-b or all). The learning rate α and the stopping rule are not stated alg. 1.
• Search-space oddities. Momentum is listed as log-uniform on [0.7, 0.99], and weight decay as log-uniform on [0.6, 0.99] tab. 1. Both are unusual ranges; the code should be checked before copying them. The meaning of "lr schedule" in [0.4, 0.9] is not stated beyond the caption tab. 1.
• Intro vs experiments on RAM. The introduction gives a RAM limit of "512Kb" §1. The experiments use 256 KB §5. The paper does not explain the difference.
• Hidden cost per evaluation. Each sample trains a DNN for up to 500 epochs tab. 1, then prunes and quantizes it and generates code to measure memory §4. The introduction says the RL overhead is negligible, but gives no timing §1. The cost of N directions grows with time and memory §5.2.
• PPO budget. The PPO hyperparameter search used 20 configurations of 20 samples each, which is not counted in the PPO-versus-ARS comparison §5.4.1. The comparison is therefore not budget-matched.
• PPO defaults outside the search space. The default n_steps is 2048 and the default ent_coef is 0.0, but the search space for n_steps is 16 to 128 and for ent_coef it starts at 1e-8 app. C tab. 5. The two columns are not on the same grid.
• Memory depends on the allocator. The paper notes that its mapping tool uses dynamic first-fit allocation, while MCUNet uses patch-based allocation. RAM comparisons therefore depend on the tool §5.4.2.
• Unstated bit widths. The search uses post-training static quantization, but the bit width used during the search is not stated. Table 3 reports int8 §4 tab. 3.
• Variance. Five seeds are run per strategy, but the text gives no numeric variance. The bands are only in the figures §5.1.
• Grey rows. In the UEA table, the grey resource values are for the highest-accuracy network, not for a network above 70% accuracy tab. 4.

Lineage

• Builds on: Augmented Random Search, "Simple random search of static linear policies is competitive for reinforcement learning" (Mania et al.) §3. ParEGO's augmented Chebyshev scalarization §4. The authors' earlier deployment pipeline §4.
• Compared against: ParEGO, TurBO, MorBO, NSGA-II and random sampling §5.1; PPO §5.4.1; MCUNet models and ProxylessNAS §5.4.2; AMC §5.4.3.
• Superseded by: not known from this wiki.