Can LLMs Revolutionize the Design of Explainable and Efficient TinyML Models?

In plain words

Small microcontrollers have very little memory and compute, so the image models that run on them must be tiny. Designing such a model by hand is slow, and automatic architecture search is usually expensive. This paper asks whether a language model can do the design work. The language model writes out a complete network as a structured description, a short list of numbers for each stage: how wide, how deep, which kernel size, and so on.

Each proposed network is checked before any training. If it is too heavy in compute or memory, the checker sends the language model a note explaining what to change. Networks that pass are trained briefly and compared on accuracy, compute and size. The ones that are not beaten on all three form a "Pareto front" of trade-offs, and that front is fed back to the language model as guidance for the next proposals.

The best networks are then trained for longer. Finally, a much larger image model acts as a teacher that shows the small network its softened answer probabilities, which helps it generalise. The paper reports that the resulting networks beat several earlier microcontroller models on CIFAR-100 while staying under the memory limit of the target chip. The same pipeline also asks the language model to explain its design choices, which the authors present as a step toward explainable design.

TL;DR

• Open-weight LLMs (Llama 3.2-3B, Llama 3.1-8B, Qwen 2.5-3B) propose five-stage CNN configurations. Candidates are filtered by MAC and SRAM estimates before training §3.1.
• A short mini-training (30 epochs) gives accuracy, MACs and parameters. Pareto dominance on these three objectives guides the LLM's next proposals §3.2.
• The chosen network gets full training (120 epochs), then logits-based distillation from a ViT-B/16 teacher for 50 epochs §4.1.3.
• Author-reported: 74.50% (LMaNet-Elite), 74.20% (LMaNet-Core), 73.00% (QwNet-Core) on CIFAR-100, under 100M MACs and 320 KB SRAM. Selection uses the CIFAR-100 test split (see Pitfalls).

Problem

Microcontrollers such as the STM32H7 have a few hundred kilobytes of SRAM. Mobile-oriented lightweight networks often need more memory than that, and hand-designed or NAS-designed microcontroller models require expensive search or manual tuning §1. Earlier LLM-for-NAS work (GENIUS, LLMatic) targets accuracy or diversity without deployment constraints such as memory and compute, and does not explain its choices §2.3. The paper assumes a hard SRAM budget of 320 KB, a MAC range of 70 to 350 million, and an STM32H743 target §4.1.4.

Method

The pipeline has three parts: an LLM-guided search over a hierarchical space §3.1, a feedback loop with Pareto selection §3.2, and ViT-based logits distillation for the final model §3.3. Figure 2 gives an overview fig. 2.

Search space §3.1 tab. 1. The network is a stack of N = 5 stages fig. 3. Each stage has L_i identical convolution blocks, each block is either a depthwise-separable block or an MBConv block, and every block in a stage shares one configuration. The stage parameters are the output channels C_out, kernel size K, stride S (1 or 2), expansion factor E, squeeze-and-excitation (SE) on or off with ratio, skip connection on or off, activation, and the block count L_i. The paper counts about 17,280 configurations per stage, hence about 17,280^5 networks in total §3.1.

Block equations §3.1. For a depthwise-separable block with input x:

\mathbf{x}' = f^{L_i}(\mathbf{x}) eq. 1

\mathbf{z} = A\big(\mathrm{BN}(\mathrm{dwconv}_{K\times K}(\mathbf{x}))\big) eq. 2

\tilde{\mathbf{z}} = \delta_{se}\,[\mathbf{z}\odot \mathrm{SEBlock}(\mathbf{z})] + (1-\delta_{se})\,\mathbf{z} eq. 3

\mathrm{SEBlock}(\mathbf{z}) = \mathbf{z}\odot\sigma\big(\mathrm{Conv}_{1\times1}(\mathrm{ReLU}(\mathrm{Conv}_{1\times1}(\mathrm{GAP}(\mathbf{z}))))\big) eq. 4

\mathbf{z}_{pw} = \mathrm{BN}\big(\mathrm{PWConv}(\tilde{\mathbf{z}})\big) eq. 5

\mathbf{x}' = \mathbf{x} + \mathbf{z}_{pw} \quad (\text{if } SC = 1),\qquad \mathbf{x}' = \mathbf{z}_{pw}\quad (\text{if } SC = 0) eq. 6

For an MBConv block, a pointwise expansion by E (with BN and activation) comes first, followed by the same depthwise-separable block, with the skip connection applied at the block level §3.1:

\mathbf{z} = A\big(\mathrm{BN}(\mathrm{PWConv}(\mathbf{x}, E))\big) eq. 7

\mathbf{z}_{pw} = \mathrm{DWSepConvBlock}(\mathbf{z}) eq. 8

Here f is one conv block, A is the activation, and δ_se ∈ {0, 1} switches SE. Equation numbers are counted in order from the LaTeX source, because the paper text does not print them in the extracted text; check them against the PDF.

Search loop §3.1, §3.2. Each iteration asks the LLM for a full five-stage JSON network, given the search space, an example, and the feedback so far. A proposal passes through these checks in order: duplicate of an earlier proposal (rejected, with feedback), MAC range (rejected, with feedback stating the current estimate), and peak SRAM estimate for an int8 model (rejected, with feedback). Only candidates that pass all checks are mini-trained.

# Feedback search loop  [§3.1] [§3.2] [code: search.py#L965-L1100]
# Input:  llm, search_space, mac_range = (70M, 350M), max_iterations
# Output: pareto_front (accuracy, MACs, params), best_candidate

def search(llm, search_space, mac_range, max_iterations):
    history = []
    pareto_front = []
    feedback = None
    for iteration in range(max_iterations):
        # The LLM proposes a whole five-stage network as JSON  [§3.1]
        config = llm.propose(search_space, feedback)
        if config in history:
            # Duplicates are not trained; the LLM is told to try again  [§3.1]
            feedback = "already evaluated; propose a unique candidate"
            continue
        history.append(config)
        # Reject if MACs fall outside the allowed range  [§3.1] [§4.1.4]
        macs = estimate_macs(config)
        if not mac_range[0] <= macs <= mac_range[1]:
            feedback = f"MAC estimate {macs} outside {mac_range}"
            continue
        # Reject if the int8 peak SRAM estimate is over the limit  [§3.1] [§4.1.4]
        if estimate_peak_sram(config) > sram_limit:
            feedback = "peak SRAM estimate too high; reduce activations"
            continue
        # Mini-train: 30 epochs on CIFAR-100  [§3.2] [§4.1.3]
        metrics = mini_train(config, epochs=30)
        # Keep the non-dominated set on accuracy, MACs and parameters  [§3.2]
        pareto_front = update_pareto(pareto_front, config, metrics)
        # Structured feedback about the front goes to the next proposal  [§3.2]
        feedback = summarise(pareto_front)
    return pareto_front

Pareto dominance §3.2. Candidate A dominates B if A is no worse on accuracy (higher is better), MACs and parameters (lower is better), and strictly better on at least one. The front keeps all non-dominated candidates. The single most accurate candidate is also tracked, even if it violates constraints §3.2.

Distillation §3.3. Teacher and student produce softened distributions at temperature T. The student loss mixes cross-entropy on labels with the KL divergence to the teacher, scaled by T²:

p_t = \mathrm{softmax}(z_t/T),\qquad p_s = \mathrm{softmax}(z_s/T) eq. 9

\mathcal{L}_{KD} = T^2\,\mathrm{KL}(p_t \,\|\, p_s) eq. 10

\mathcal{L} = \alpha\,\mathcal{L}_{CE} + (1-\alpha)\,\mathcal{L}_{KD} eq. 11

\alpha \leftarrow \alpha + (\alpha_{final}-\alpha)\,\frac{\text{epoch}}{\text{num\_epochs}} eq. 12

# Logits distillation with a per-epoch alpha update  [§3.3] [eq. 9]-[eq. 12]
# Input:  student_logits (B, C), teacher_logits (B, C), labels (B,)
# Output: scalar loss

def distillation_loss(student_logits, teacher_logits, labels, alpha, T):
    # Soften both output distributions with temperature T  [eq. 9]
    p_teacher = softmax(teacher_logits / T)
    log_p_student = log_softmax(student_logits / T)
    # KL(teacher || student), scaled by T squared  [eq. 10]
    kd = T**2 * kl_divergence(p_teacher, log_p_student)
    ce = cross_entropy(student_logits, labels)
    # alpha weights the label loss, (1 - alpha) the teacher loss  [eq. 11]
    return alpha * ce + (1 - alpha) * kd

Hyperparameters

Results

Author-reported. Setup: CIFAR-100 at 160 × 160, STM32H743 target, baseline numbers as listed in the paper tab. 2. The paper does not say whether baselines were retrained under its recipe.

Peak SRAM, from the text §4.2.2 and the figure fig. 5: LMaNet-Elite 165 KB, LMaNet-Core 215 KB, QwNet-Core 227 KB, MCUNet-in3 293 KB, XiNet-Class 204 KB, MCUNet-in4 456 KB.

Effect of ViT distillation tab. 4: LMaNet-Core 72.88 → 74.20 (+1.32), QwNet-Core 71.60 → 73.00 (+1.40), LMaNet-Elite 73.15 → 74.50 (+1.35).

Mini-phase candidates, 30 epochs, search statistics tab. 3: Llama 3.2-3B generated 127 candidates (27 on the Pareto front), Llama 3.1-8B 100 (29), Qwen 2.5-3B 213 (43). Average Pareto accuracy was 60.50, 61.50 and 61.90 respectively.

Paper vs code

Official repository: christophezei/llm-nas-kd-explainability, linked from the abstract [abstract]. The owner matches the first author's name, and the README cites the same arXiv id. Commit 778742a8. The repository has no LICENSE file, so its code is not licensed for reuse; cite by path and line only. Its README mentions an IJCNN 2025 award, so it is probably later than arXiv v1.

• SRAM limit in search. Paper: 320 KB is a hard constraint §4.1.4, §5.1. Code: candidates are rejected above 0.512 in MiB units, about 512 KB search.py helper.py.
• Selection on test accuracy. Paper: no explicit statement. Code: Pareto dominance, the best-candidate choice, and the LLM feedback all use test_acc search.py search.py search.py. The validation split is used only for the "promising" gate search.py.
• Promising gate. Not in paper. A candidate is "promising" if validation accuracy exceeds 70%, or validation F1 exceeds 0.55, or the accuracy slope exceeds 0.5, or the average improvement exceeds 1.0 search.py. Only promising candidates can become best search.py.
• Mini-phase warm-up. Paper: 10-epoch warm-up §4.1.3. Code: warm-up is set but unused; the scheduler is a step schedule search.py.
• Full-phase learning rate. Paper: 0.5 §4.1.3. Shipped training script: 0.3 train.py.
• Full-phase extras not in paper. Label smoothing 0.1 train.py; AutoAugment with the CIFAR-10 policy train.py; RandomAffine (±10°, translate 0.1, scale 0.8–1.2), applied to MixUp images train.py train.py; dropout 0.05 train.py; 5% of training images held out for validation train.py.
• No checkpoint selection. Final-epoch weights are saved; the best_val_acc value is never updated and is reported as zero train.py train.py.
• Teacher checkpoint. Paper: pretrained ViT-B/16, footnote points to the ImageNet-21k base model §4.1.3. Code: timm ViT-B/16 with a 100-way head, loaded from a checkpoint whose name says it was fine-tuned on CIFAR-100 kd_phase.py. The paper's 93% teacher accuracy §5.2 is not tied to a checkpoint.
• KD data and optimiser. Paper: none of the KD learning rate, optimiser or batch size is stated. Code: 80% of training images, stratified split with seed 42 kd_phase.py; SGD with Nesterov 0.9, weight decay 1e-4, OneCycle schedule with peak 2e-3, batch 256 kd_phase.py kd_phase.py kd_phase.py; label smoothing 0.1 in the cross-entropy term kd_phase.py. Teacher input is 224 px with mean and std 0.5 kd_phase.py kd_phase.py.
• Alpha direction. Paper prose: cross-entropy dominates early and the distillation weight rises over training §3.3. Equation and code: the cross-entropy weight is α and rises from 0.4 to 0.8, so the distillation weight falls eq. 11 kd_phase.py. The code updates α recursively each epoch kd_phase.py.
• Residual connections. Paper: skip connection is x + z_pw eq. 6. Code: applied only when stride is 1 model.py; a 1×1 conv and BN projection is added when channels differ model.py.
• Stride placement. Paper: not stated. Code: stride applies to the first layer of each block; later layers in the block use stride 1 model.py.
• Expansion factor for depthwise blocks. Paper: E is in the search space for every block tab. 1. Code: E is not used for depthwise-separable blocks model.py.
• Network head. Paper: no stem or head described. Code: blocks, global average pool, dropout, linear classifier; no stem convolution model.py.
• MBConv expansion. Paper: expansion before the depthwise-separable block eq. 7. Code: pointwise expansion, depthwise convolution, SE, pointwise projection in one sequence model.py. The two match in effect.
• SRAM estimator and size. Paper: not described. Code: int8 peak estimate from a forward-hook sum of input size plus one module's output at a time helper.py helper.py. Parameter size is counted as 4 bytes per parameter helper.py.
• MACs tool. Paper: not named. Code: ptflops get_model_complexity_info helper.py.
• LLM explanations. Paper: explanations "guide further iterations" fig. 2. Code: explanations are printed for each valid candidate, not stored or fed back search.py.
• Search prompt. Paper figure: targets 70% and 350M MACs fig. 4. Code prompt: targets 75% against a stated 75% state of the art at 259M MACs search.py. Code's example answer uses expansion factor 2 and SE ratio 0, outside the listed sets search.py.
• LLM choice. Paper: three LLMs compared §5.1. Code: only Llama 3.1-8B is active; the 3B and Qwen lines are commented out search.py.
• Shipped configurations. The three JSON files contain only candidate_config, with no metrics key LMaNet_elite.json. load_best_candidate reads metrics and would raise a KeyError helper.py. So train.py and model.py main() fail on these files as shipped.
• Explanation vs shipped model. Figure 6 says strides are 1 in later blocks fig. 6. The shipped LMaNet-Elite ends with a stride-2 block and uses Leaky ReLU only in block 3 LMaNet_elite.json.

Pitfalls

• Selection on the test split. The search ranks, picks and feeds back test accuracy search.py. The headline 74.50% is therefore not a held-out score. For any reuse, select on a validation split and report test once.
• SRAM budget is loose in code. The search admits models up to about 512 KB while the paper says 320 KB. Reported SRAM figures §4.2.2 come from an estimator the paper does not describe; the repository does not reproduce them, since the shipped configurations have no stored metrics.
• The SRAM estimator is crude. It counts one module's output at a time helper.py. It does not keep residual inputs or expanded tensors alive together, so it can under- or over-state real runtime memory. Use a real microcontroller runtime to confirm fit.
• Mini-phase warm-up. Implement the code's step schedule if you reproduce the search; the paper's warm-up is not in the code.
• Full-phase learning rate. Paper 0.5, shipped code 0.3. Pick one and say which.
• Alpha schedule. Paper prose and the equation disagree on direction. Eq. 11 with α from 0.4 to 0.8 gives the cross-entropy term more weight over time, the reverse of the prose. Implement the equations as written and note the conflict.
• α update form. Eq. 12 has α on both sides. The code updates the running value each epoch kd_phase.py. A closed-form reading gives different values.
• Teacher identity changes the result. The paper says pretrained ViT-B/16 and cites the ImageNet-21k base model. The code uses a CIFAR-100-tuned teacher. The distillation gain tab. 4 may depend on that choice.
• Where stride and skip apply. The paper leaves these unstated. The code's choices (stride on the first layer only, skip only at stride 1, projection when channels change) change the architecture and the parameter count.
• Search budget. Table 3 counts 127, 100 and 213 candidates tab. 3; the code runs 500 iterations, and rejected or duplicate proposals are not stored. Counts of "candidates" are not iterations.
• Baselines are quoted, not re-run. MCUNet and XiNet figures come from the cited works tab. 2. Their recipes differ. Without distillation, the proposed models are at 72.88 (LMaNet-Core), 71.60 (QwNet-Core) and 73.15 (LMaNet-Elite) tab. 4. So QwNet-Core without distillation is below MCUNet-in4's 72.86, and the margins over the baselines are small.
• Speed claim. The text says all three models are "over 4 times faster" than MCUNet's 12.5 days. Against QwNet-Core's 3.5 days that is about 3.6 times tab. 2.
• Model-size wording. The text gives LMaNet-Core as 0.51 MB §4.2.1, which is its parameter count; its flash size in Table II is 0.53 MB tab. 2.
• Flash size definition. Table II flash values are not defined in the paper; code reports float32 parameter size instead helper.py.
• Reproducing the search costs LLM inference. The code loads an 8B model with a 16384-token context and no 4-bit quantisation search.py. Budget GPU memory and time accordingly.
• Licensing. The repository has no LICENSE file, and it ships trained weights (.pth) with no licence either. Do not copy code or weights into other work without permission; cite the paper's equations and the code by path and line.

Lineage

• Builds on: knowledge distillation with logits [§3.3, Hinton et al. cited as hinton2015distilling]; MobileNet-style MBConv blocks §3.1.
• Compared against: MCUNet-in3 and MCUNet-in4 tab. 2, XiNet-Class tab. 2.
• Related LLM-for-NAS work: GENIUS and LLMatic §2.3. The paper gives no arXiv ids for them, so none are linked here.
• Superseded by: not known from this wiki.