🤖 machine learning

LLMForge: Multi-Backend Hardware-Aware Neural Architecture Search with Infinite-Head Attention for Edge Language Models

LLMForge is a hardware-aware neural architecture search framework that leverages Infinite-Head Attention and a multi-backend cost model to automatically generate optimized sub-billion-parameter language models tailored to specific edge device constraints, achieving significant improvements in energy efficiency, latency, and model accuracy compared to existing baselines.

Original authors: Xinting Jiang, Junyi Luo, Ruichen Qi, Kauna Lei, Ben Laurie, Gregory Kielian, Mehdi Saligane

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xinting Jiang, Junyi Luo, Ruichen Qi, Kauna Lei, Ben Laurie, Gregory Kielian, Mehdi Saligane

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build the perfect, tiny robot brain that can run on a smartphone or a smartwatch. You want it to be smart enough to understand language, but small enough to fit in your pocket without draining the battery or making the phone hot.

The problem is that the standard "blueprints" for these brains (called Transformers) were designed for massive supercomputers. When you try to shrink them down for edge devices, they often get stuck in traffic jams, run out of memory, or eat too much energy.

LLMForge is a new tool created by researchers to solve this. Think of it as a super-smart, hardware-aware architect that doesn't just shrink the brain; it completely redesigns the blueprint based on the specific "terrain" of the device it will live on.

Here is how LLMForge works, broken down into three simple parts:

1. The Infinite-Head Attention (IHA): Breaking the Rules

Traditional robot brains have a rigid rulebook. If you want to increase the number of "thinking heads" (processors that look at different parts of a sentence), you are forced to make each head smaller. It's like a pizza: if you cut it into more slices, every slice gets smaller. This limits how smart the brain can be.

LLMForge's innovation: It throws out the rigid rulebook. With Infinite-Head Attention, the architect can change the number of slices, the size of the slices, and the type of toppings independently for every single layer of the brain.

The Analogy: Imagine a construction crew that usually has to build rooms in identical sizes. LLMForge gives them a magic tool that lets them build a tiny closet next to a massive ballroom, then a medium office, all within the same floor plan. This expands the number of possible blueprints by 400 times, allowing them to find a shape that fits perfectly into a specific device's constraints.

2. Forge-Former: The Crystal Ball

Building and testing a new brain blueprint from scratch is expensive and slow. It's like baking a cake just to see if the recipe works, then throwing it away. If you have thousands of recipes to test, you'd go bankrupt.

LLMForge's solution: They built a Crystal Ball called Forge-Former.

How it works: Instead of baking every single cake, the Crystal Ball looks at the recipe (the blueprint) and predicts exactly how good the cake will taste (how smart the model will be) and how much energy it will use.
The Result: This crystal ball is much more accurate than previous "guessing" tools. It allows the system to test thousands of designs in the time it used to take to test just a few, saving massive amounts of time and energy.

3. Forge-DSE: The Multi-Tool Navigator

Different devices have different problems. A high-end graphics card (GPU) might be limited by how fast it can move data, while a tiny chip in a smartwatch might be limited by how much heat it can handle. A design that is perfect for one might be terrible for another.

LLMForge's approach: The Forge-DSE engine acts like a navigator with a multi-tool.

It doesn't just look for the "smartest" model. It looks for the best trade-off (a "Pareto front") between being smart, fast, and energy-efficient.
It tests the designs against four different types of hardware (like a high-speed GPU, a specialized chip, and a ring-shaped chip).
The Outcome: The system realizes that "one size does not fit all." It produces different, unique blueprints for each device.
- For a speed-focused device, it builds a brain that is wide and shallow.
- For an energy-focused device, it builds a brain that is deep and narrow.

The Results: Real-World Wins

The researchers tested this system on two sizes of models (around 100 million and 300 million parameters) and compared them to existing popular models like SmolLM2 and Qwen.

The "Accuracy" Model: One of their new designs was smarter (lower error rate) than the competition, even though it had fewer parameters (it was smaller).
The "Energy" Model: Another design used 40% less energy per word generated compared to the standard models.
The "Speed" Model: A third design was 43% faster at starting to talk (Time-to-First-Token) and finishing sentences.

Summary

LLMForge is a system that stops trying to force a "one-size-fits-all" brain onto every device. Instead, it uses a flexible design language, a fast prediction crystal ball, and a smart navigator to engineer a custom brain for every specific piece of hardware. The result is smaller, faster, and more energy-efficient AI that actually fits in your pocket.

Technical Summary: LLMForge

Problem Statement

Sub-billion-parameter Transformer language models are increasingly deployed on edge devices to leverage privacy, latency, and cost advantages. However, on-device inference is constrained by tight memory bandwidth, energy, and thermal budgets. Existing approaches often rely on post-training compression (quantization, pruning), which leaves the underlying architectural degrees of freedom unexplored. Furthermore, current Hardware-Aware Neural Architecture Search (NAS) pipelines for Transformers suffer from three limitations:

Restricted Search Space: Attention parameterizations often fix per-head dimensions and tie key/value capacities, limiting the feasible shape configurations to a single head-count choice.
Evaluation Bottleneck: Evaluating candidates by training from scratch is computationally expensive.
Coarse Cost Modeling: Hardware costs are frequently modeled via single accelerator classes or coarse proxies (e.g., parameter count), masking substrate-dependent re-ranking of candidates.

Methodology

The authors present LLMForge, a hardware-aware NAS framework that integrates three composable contributions into a single NSGA-II evolutionary search loop.

1. Infinite-Head Attention (IHA)

IHA decouples the attention shape parameters ( $n_h$ , $n_{kv}$ , $d_{qk}$ , $d_v$ ) as independent per-layer variables. Unlike standard Multi-Head Attention (MHA) or Grouped-Query Attention (GQA), which enforce divisibility constraints ( $d_{model} = n_h \cdot d_h$ ) and Q/K–V coupling ( $d_h = d_{qk} = d_v$ ), IHA relaxes these constraints. It only preserves the GQA grouping condition ( $n_{kv} | n_h$ ).

Impact: This expands the feasible per-layer attention configuration space by approximately 400× (from 27 configurations under GQA to 11,250 under IHA within the defined search ranges).

2. Forge-Former (Accuracy Surrogate)

To address the evaluation bottleneck, the authors introduce Forge-Former, a Transformer-encoder-based surrogate model.

Architecture: It encodes an architecture as a sequence of layer tokens, where each token concatenates the layer's IHA parameters. Scalar fields are lifted via linear layers and summed with positional embeddings. A pre-LN Transformer encoder captures within-layer interactions (via feedforward sublayers) and cross-layer interactions (via self-attention).
Performance: On the IHA test split, Forge-Former achieves a Spearman $\rho = 0.75$ and Kendall $\tau = 0.58$ , outperforming MLP and random-forest baselines by 1.4 to 2.5×. It enables NSGA-II selection at a fraction of the cost of training from scratch.

3. Forge-DSE (Design-Space-Exploration Engine)

Forge-DSE is the multi-objective search engine that pairs Forge-Former with a multi-backend hardware cost model.

Search Loop: It evolves populations of IHA-parameterized architectures using NSGA-II with non-dominated sorting and crowding-distance survival. It employs custom mutation operators (deletion, duplication, rotation, numerical perturbation) and a feasibility repair step.
Surrogate Co-Evolution: The framework optionally refines Forge-Former during the search by real-training informative samples from the current population and fine-tuning the surrogate via experience replay, correcting for distribution shift.
Hardware Backends: The system supports three distinct backends:
1. Backend A: Measured costs on an NVIDIA A100 GPU using ZEUS.
2. Backend B: Analytical costs via Timeloop for systolic (Gemmini), row-stationary (Eyeriss), and fused-attention (FLAT) dataflows.
3. Backend C: A multi-chip ring-dataflow simulator (rDXE) for edge accelerators, performing joint architecture-and-chip co-search.

Key Results

Surrogate Validation

On the IHA dataset, Forge-Former significantly outperforms baselines in ranking fidelity and top-X% recovery. It reduces the Mean Absolute Error (MAE) for the top 5% of architectures by approximately 5× compared to MLP and Random Forest baselines. t-SNE analysis confirms that Forge-Former embeddings better reflect performance similarity, clustering high-loss tails tightly.

Multi-Backend Search Performance

Across four hardware substrates (A100, Gemmini, Eyeriss, FLAT), the searches converge to distinct Pareto-front architectures that track each substrate's specific cost bottlenecks:

A100 (Bandwidth-rich): Results in flat MLP widths and the widest attention heads.
Eyeriss/Gemmini (KV-cache bound): Collapse $n_{kv}$ into aggressive grouped-query attention and compensate with wider MLP capacity.
FLAT (Flexible): Produces the smallest model with shallow depth and narrow heads.
Improvements: Compared to the SmolLM2-135M baseline, the search yields reductions of 5–7% in validation loss, 16–63% in energy per token, and 16–59% in latency metrics (TTFT/TPOT).

Scaled Validation on rDXE

The authors performed joint architecture-and-chip co-search on the rDXE multi-chip ring substrate, retraining selected candidates on FineWeb-Edu-10BT (13.1B tokens) against SmolLM2 and Qwen baselines:

Accuracy Variant (347M params): Achieved the lowest validation loss (2.798) and competitive benchmark performance with fewer parameters than SmolLM2-360M and Qwen-0.5B.
Energy Variant (294M params): Reduced energy per token by 40%.
Latency Variant (365M params): Reduced Time-to-First-Token (TTFT) and Time-Per-Output-Token (TPOT) by 43%.

Significance and Claims

The paper claims that LLMForge demonstrates that hardware-aware NAS can produce substrate-conditioned Pareto fronts rather than a single "best" architecture. By decoupling attention parameters via IHA and utilizing a high-fidelity encoder-based surrogate, the framework effectively reshapes per-layer width and depth profiles to match specific hardware bottlenecks.

The authors position LLMForge as a methodological contribution that enables the discovery of deployment-aware architecture-and-chip configuration pairs. They note that while the energy and latency reductions are modest, they primarily benefit the operating cost and carbon profile of repeated edge inference. The work is limited to sub-billion-parameter scales ( $\le$ 500M) due to the computational cost of generating the training corpus for the surrogate, with scaling to billion-parameter models identified as a future step.