🤖 machine learning

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

LLMForge is een hardwarebewust raamwerk voor neurale architectuurzoekopdrachten dat gebruikmaakt van Infinite-Head Attention en een multi-backend kostenmodel om automatisch geoptimaliseerde taalmodellen met minder dan een miljard parameters te genereren die zijn afgestemd op specifieke randapparatuurbeperkingen, met aanzienlijke verbeteringen in energie-efficiëntie, latentie en modelnauwkeurigheid ten opzichte van bestaande basismodellen.

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

Gepubliceerd 2026-05-19

📖 4 min leestijd☕ Koffiepauze-leesvoer

CC BY 4.0

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

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je probeert de perfecte, kleine robotbrein te bouwen dat kan draaien op een smartphone of smartwatch. Je wilt dat het slim genoeg is om taal te begrijpen, maar klein genoeg om in je zak te passen zonder de batterij te leegtrekken of de telefoon heet te maken.

Het probleem is dat de standaard "blauwdrukken" voor deze breinen (zogenaamde Transformers) zijn ontworpen voor enorme supercomputers. Als je ze probeert te verkleinen voor randapparaten, raken ze vaak vast in file, raken ze het geheugen op, of verbruiken ze te veel energie.

LLMForge is een nieuw hulpmiddel dat door onderzoekers is ontwikkeld om dit op te lossen. Denk er als een super-slimme, hardware-bewuste architect die niet alleen het brein verkleint; het herbouwt de blauwdruk volledig op basis van het specifieke "terrein" van het apparaat waarop het zal leven.

Hier is hoe LLMForge werkt, opgesplitst in drie eenvoudige onderdelen:

1. De Oneindige-Kop Aandacht (IHA): De Regels Breken

Traditionele robotbreinen hebben een rigide regelboek. Als je het aantal "denkkoppen" (processors die naar verschillende delen van een zin kijken) wilt verhogen, word je gedwongen elke kop kleiner te maken. Het is als een pizza: als je hem in meer stukken snijdt, wordt elk stuk kleiner. Dit beperkt hoe slim het brein kan zijn.

De innovatie van LLMForge: Het gooit het rigide regelboek weg. Met Oneindige-Kop Aandacht kan de architect het aantal stukken, de grootte van de stukken en het type topping onafhankelijk van elkaar aanpassen voor elke enkele laag van het brein.

De Analogie: Stel je een bouwteam voor dat meestal kamers in identieke maten moet bouwen. LLMForge geeft hen een magisch gereedschap dat hen in staat stelt een kleine kast naast een enorme balzaal te bouwen, gevolgd door een middelgroot kantoor, allemaal binnen hetzelfde plattegrond. Dit vergroot het aantal mogelijke blauwdrukken met 400 keer, waardoor ze een vorm kunnen vinden die perfect past binnen de beperkingen van een specifiek apparaat.

2. Forge-Former: De Kristallen Bol

Het bouwen en testen van een nieuwe breinblauwdruk vanaf nul is duur en traag. Het is alsof je een taart bakt om te zien of het recept werkt, en hem vervolgens weggooit. Als je duizenden recepten moet testen, ga je failliet.

De oplossing van LLMForge: Ze bouwden een Kristallen Bol genaamd Forge-Former.

Hoe het werkt: In plaats van elke taart te bakken, kijkt de Kristallen Bol naar het recept (de blauwdruk) en voorspelt precies hoe goed de taart zal smaken (hoe slim het model zal zijn) en hoeveel energie het zal verbruiken.
Het Resultaat: Deze kristallen bol is veel nauwkeuriger dan eerdere "gok"-hulpmiddelen. Het stelt het systeem in staat duizenden ontwerpen te testen in de tijd die voorheen nodig was om er slechts een paar te testen, wat enorme hoeveelheden tijd en energie bespaart.

3. Forge-DSE: De Multi-Tool Navigator

Verschillende apparaten hebben verschillende problemen. Een high-end grafische kaart (GPU) kan beperkt zijn door hoe snel het data kan verplaatsen, terwijl een kleine chip in een smartwatch beperkt kan zijn door hoeveel hitte het kan verdragen. Een ontwerp dat perfect is voor het ene, kan vreselijk zijn voor het andere.

De aanpak van LLMForge: De Forge-DSE-engine fungeert als een navigator met een multi-tool.

Het zoekt niet alleen naar het "slimste" model. Het zoekt naar de beste afweging (een "Pareto-front") tussen slim zijn, snel zijn en energiezuinig zijn.
Het test de ontwerpen tegen vier verschillende soorten hardware (zoals een high-speed GPU, een gespecialiseerde chip en een ringvormige chip).
Het Resultaat: Het systeem beseft dat "één maat niet voor iedereen past". Het produceert verschillende, unieke blauwdrukken voor elk apparaat.
- Voor een snelheidsgerichte apparatuur bouwt het een brein dat breed en ondiep is.
- Voor een energiegerichte apparatuur bouwt het een brein dat diep en smal is.

De Resultaten: Wereldwijde Overwinningen

De onderzoekers testten dit systeem op twee maten van modellen (rond de 100 miljoen en 300 miljoen parameters) en vergeleken ze met bestaande populaire modellen zoals SmolLM2 en Qwen.

Het "Nauwkeurigheids"-Model: Een van hun nieuwe ontwerpen was slimmer (lagere foutmarge) dan de concurrentie, zelfs al had het minder parameters (het was kleiner).
Het "Energie"-Model: Een ander ontwerp verbruikte 40% minder energie per gegenereerd woord in vergelijking met de standaardmodellen.
Het "Snelheids"-Model: Een derde ontwerp was 43% sneller in het starten van het spreken (Time-to-First-Token) en het voltooien van zinnen.

Samenvatting

LLMForge is een systeem dat stopt met proberen een "één-maat-voor-alles"-brein op elk apparaat te forceren. In plaats daarvan gebruikt het een flexibel ontwerptaal, een snelle voorspellende kristallen bol en een slimme navigator om een op maat gemaakt brein te bouwen voor elk specifiek stuk hardware. Het resultaat is een kleinere, snellere en energiezuinigere AI die daadwerkelijk in je zak past.

Technische Samenvatting: LLMForge

Probleemstelling

Sub-miljard-parameter Transformer-taalmodellen worden steeds vaker op randapparaten ingezet om voordelen op het gebied van privacy, latentie en kosten te benutten. Echter, inferentie op het apparaat wordt beperkt door krappe bandbreedte, energie- en thermische budgetten. Bestaande benaderingen vertrouwen vaak op compressie na training (kwantisatie, pruning), waardoor de onderliggende architecturale vrijheidsgraden onbenut blijven. Bovendien lijden huidige Hardware-Aware Neural Architecture Search (NAS)-piplines voor Transformers aan drie beperkingen:

Beperkte Zoekruimte: Attention-parametriseringen fixeren vaak per-head-dimensies en koppelen key/value-capaciteiten, waardoor de haalbare vormconfiguraties worden beperkt tot één keuze voor het aantal heads.
Evaluatieknelpunt: Het evalueren van kandidaten door ze vanaf nul te trainen is computatievermogenintensief.
Grove Kostenmodellering: Hardwarekosten worden frequent gemodelleerd via enkele acceleratorklassen of grove proxies (bijv. aantal parameters), waardoor een substraatafhankelijke herordening van kandidaten wordt gemaskeerd.

Methodologie

De auteurs presenteren LLMForge, een hardware-aware NAS-framework dat drie composable bijdragen integreert in één enkele NSGA-II evolutionaire zoeklus.

1. Infinite-Head Attention (IHA)

IHA ontkoppelt de attention-vormparameters ( $n_h$ , $n_{kv}$ , $d_{qk}$ , $d_v$ ) als onafhankelijke per-laagvariabelen. In tegenstelling tot standaard Multi-Head Attention (MHA) of Grouped-Query Attention (GQA), die delbaarheidsbeperkingen afdwingen ( $d_{model} = n_h \cdot d_h$ ) en Q/K–V-koppeling ( $d_h = d_{qk} = d_v$ ), versoepelt IHA deze beperkingen. Het behoudt alleen de GQA-groeperingsvoorwaarde ( $n_{kv} | n_h$ ).

Impact: Dit breidt de haalbare per-laag attention-configuratieruimte uit met ongeveer 400× (van 27 configuraties onder GQA naar 11.250 onder IHA binnen de gedefinieerde zoekbereiken).

2. Forge-Former (Accuracy Surrogate)

Om het evaluatieknelpunt aan te pakken, introduceren de auteurs Forge-Former, een op Transformer-encoder gebaseerd surrogate-model.

Architectuur: Het encodeert een architectuur als een sequentie van laag-tokens, waarbij elk token de IHA-parameters van de laag concateneert. Scalar-velden worden via lineaire lagen gelift en opgeteld bij positionele embeddings. Een pre-LN Transformer-encoder vangt interacties binnen de laag (via feedforward-sublagen) en interacties tussen lagen (via self-attention) op.
Prestaties: Op de IHA-testsplit bereikt Forge-Former een Spearman $\rho = 0,75$ en Kendall $\tau = 0,58$ , en presteert het 1,4 tot 2,5× beter dan MLP- en random-forest-baselines. Het stelt NSGA-II-selectie mogelijk tegen een fractie van de kosten van training vanaf nul.

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

Forge-DSE is de multi-objectieve zoekengine die Forge-Former koppelt aan een multi-backend hardwarekostenmodel.

Zoeklus: Het evolueert populaties van IHA-geparameerde architecturen met NSGA-II, gebruikmakend van non-dominated sorting en overleving op basis van crowding-distance. Het maakt gebruik van aangepaste mutatie-operatoren (verwijdering, duplicatie, rotatie, numerieke perturbatie) en een herstelstap voor haalbaarheid.
Surrogate Co-evolutie: Het framework verfijnt Forge-Former optioneel tijdens de zoektocht door het opnieuw trainen van informatieve samples uit de huidige populatie en het fine-tunen van het surrogate via experience replay, waardoor verschuivingen in verdeling worden gecorrigeerd.
Hardware Backends: Het systeem ondersteunt drie distincte backends:
1. Backend A: Gemeten kosten op een NVIDIA A100 GPU met ZEUS.
2. Backend B: Analytische kosten via Timeloop voor systolische (Gemmini), rij-gestationeerde (Eyeriss) en gefuseerde-attention (FLAT) dataflows.
3. Backend C: Een multi-chip ring-dataflow-simulator (rDXE) voor randaccelerators, die een gezamenlijke architectuur-en-chip co-zoektocht uitvoert.

Belangrijkste Resultaten

Surrogate Validatie

Op de IHA-dataset presteert Forge-Former significant beter dan baselines in rangorde-trouw en herstel van top-X%. Het verlaagt de Mean Absolute Error (MAE) voor de top 5% van architecturen met ongeveer 5× ten opzichte van MLP- en Random Forest-baselines. t-SNE-analyse bevestigt dat Forge-Former-embeddings prestatie-相似heid beter weerspiegelen, waarbij de hoge-verliesstaarten strak worden geclusterd.

Multi-Backend Zoekprestaties

Over vier hardware-substraten (A100, Gemmini, Eyeriss, FLAT) convergeren de zoektochten naar distincte Pareto-front-architecturen die de specifieke kostenknelpunten van elk substraat volgen:

A100 (Rijk aan bandbreedte): Resulteert in vlakke MLP-breedtes en de breedste attention-heads.
Eyeriss/Gemmini (KV-cache-gebonden): Laten $n_{kv}$ instorten tot agressieve grouped-query attention en compenseren met bredere MLP-capaciteit.
FLAT (Flexibel): Produceert het kleinste model met ondiepe diepte en smalle heads.
Verbeteringen: Ten opzichte van de SmolLM2-135M-baseline levert de zoektocht reducties op van 5–7% in validatieverlies, 16–63% in energie per token, en 16–59% in latentie-metrics (TTFT/TPOT).

Geschaalde Validatie op rDXE

De auteurs voerden een gezamenlijke architectuur-en-chip co-zoektocht uit op het rDXE multi-chip ring-substraat, waarbij geselecteerde kandidaten opnieuw werden getraind op FineWeb-Edu-10BT (13,1B tokens) tegenover SmolLM2- en Qwen-baselines:

Accuracy Variant (347M params): Bereikte het laagste validatieverlies (2,798) en concurrerende benchmarkprestaties met minder parameters dan SmolLM2-360M en Qwen-0.5B.
Energy Variant (294M params): Verminderde energie per token met 40%.
Latency Variant (365M params): Verminderde Time-to-First-Token (TTFT) en Time-Per-Output-Token (TPOT) met 43%.

Betekenis en Claims

Het artikel claimt dat LLMForge aantoont dat hardware-aware NAS substraat-geconditioneerde Pareto-fronts kan produceren in plaats van één enkele "beste" architectuur. Door attention-parameters te ontkoppelen via IHA en gebruik te maken van een hoog-trouw encoder-gebaseerd surrogate, herschikt het framework effectief per-laag breedte- en diepteprofielen om te matchen met specifieke hardwareknelpunten.

De auteurs positioneren LLMForge als een methodologische bijdrage die de ontdekking mogelijk maakt van deployment-bewuste architectuur-en-chip configuratieparen. Zij merken op dat hoewel de energie- en latentiereducties bescheiden zijn, deze primair ten goede komen aan de operationele kosten en het koolstofprofiel van herhaalde inferentie op randapparaten. Het werk is beperkt tot sub-miljard-parameter schalen ( $\le$ 500M) vanwege de computatiekosten van het genereren van het trainingscorpus voor het surrogate, met schaalvergroting naar miljard-parameter modellen geïdentificeerd als een toekomstige stap.