Het artikel beschrijft de technische realisatie van een generatief AI-beeldmodel op een RP2350-microcontroller (gebruikt in de Raspberry Pi Pico 2). De auteur heeft een latent flow diffusion transformer (DiT) en een variational auto-encoder (VAE) geïmplementeerd, getraind op de FFHQ-dataset.
Belangrijkste technische details:
- Hardware: Gebruik van een Waveshare RP2350-Plus board met een Pimoroni VGA-board voor visuele output.
- Model: Twee varianten met 1,7 en 2,9 miljoen parameters, gecomprimeerd tot minder dan 4 MB via
int8 kwantisering.
- Proces: Beelden worden gegenereerd in een latente ruimte (16x16x8) en vervolgens gedecodeerd naar 128x128 RGB-pixels.
- Optimalisaties: Toepassing van Classifier Free Guidance (CFG), weight streaming via DMA vanuit het flashgeheugen, en overclocking van de Cortex-M33 cores naar 300 MHz.
- Prestaties: De generatie van een beeld neemt 10 tot 20 seconden in beslag.
Het project demonstreert dat complexe AI-modellen, mits sterk geoptimaliseerd, kunnen draaien op zeer beperkte hardware.
AI-beeldgeneratie op een RP2350-microcontroller
Onlangs ben ik naar het project teruggekeerd en ik ben erin geslaagd een generatief AI-beeldmodel te implementeren op een RP2350-microcontroller van 1 dollar, zoals gebruikt in de Raspberry Pi Pico 2. Het model en de inference-code zijn kleiner dan 4 MB en draaien op de dual-core Cortex-M33 met 520 KB RAM.
Wat doet het?
De hardware bestaat uit een RP2350-microcontroller op een Waveshare RP2350-Plus board, aangesloten op een Pimoroni VGA-board. Het VGA-board voegt in feite een resistor-gebaseerde DAC en een VGA-connector toe, waardoor aansluiting op een VGA-monitor mogelijk is. De monitor is niet strikt noodzakelijk, aangezien het gegenereerde beeld ook via de USB-verbinding kan worden uitgelezen. HDMI-ondersteuning zou evenzeer mogelijk zijn, maar ik had geen passend board bij de hand.
Het apparaat genereert en toont RGB-afbeeldingen van menselijke gezichten met een resolutie van 128×128 pixels, waarbij elke afbeelding 10 tot 20 seconden in beslag neemt. Deze worden vervolgens weergegeven op een VGA-monitor of gestreamd via USB.
Het model implementeert een latent flow diffusion transformer (DiT), vergelijkbaar met wat wordt gebruikt in diffusiemodellen zoals Flux. Er zijn twee varianten met respectievelijk 2,9 en 1,7 miljoen parameters, wat ongeveer 4000 keer minder is dan Flux. Het ondersteunt conditionele generatie in vijf klassen: vier combinaties van geslacht en glimlach, plus een ongeconditioneerde klasse. Het model is getraind op de FFHQ-dataset.
Het is verbazingwekkend dat een model van deze omvang in staat is om complexe beelden te genereren. Veel MNIST-speelgoedprojecten voor diffusie gebruiken veel meer parameters en zijn nauwelijks in staat om iets coherents te genereren.
Hoe werkt het?
Hoewel Claude Code met Fable 5 een groot deel van het zware werk bij de implementatie van de code heeft gedaan, nam de ontwikkeling bijna twee weken in beslag, inclusief vele ablaties en nachtelijke trainingsruns op een RTX 5090. Wat me opviel, is dat de meeste optimalisaties die hielpen bij grote modellen ook noodzakelijk waren om dit micro-model werkend te krijgen.
Latente Diffusie
Het model genereert geen beelden in de pixelruimte, maar genereert een gecomprimeerde latente representatie van het beeld. Deze aanpak werd gepionierd door Stable Diffusion (Rombach et al.) en is nu in feite de standaard voor diffusiemodellen. Voor dit model wordt de pixelruimte van 128×128×3 met een factor 24 gecomprimeerd naar een 16×16×8 representatie met behulp van een variational auto-encoder (VAE).
De VAE is afzonderlijk getraind op de FFHQ-dataset. In een ongebruikelijke opstelling is de encoder veel groter dan de decoder. De reden hiervoor was om de grote encoder een latente representatie te laten leren die de decodering optimaliseert voor een kleine, on-device decoder. Voor de inference op de microcontroller is alleen de decoder nodig. LPIPS werd gebruikt als visuele perceptuele verliesfunctie.
Er zijn twee VAE's getraind: één met 115k parameters en één met 494k parameters, respectievelijk voor het 'snelle' en het 'kwaliteitsmodel'. De decoder met 494k parameters is slechts marginaal beter (bijvoorbeeld in de weergave van tanden), wat aantoont dat de latente dimensie de beperkende factor is.
Flow-based Diffusiemodel
Diffusiemodellen worden vaak verpakt in onnodig complexe wiskundige taal (DDPM, DDIM...), terwijl ze relatief eenvoudig zijn. Ik heb een flow-matching objectief gebruikt om de diffusie-transformer te trainen.
In essentie trainen we een model dat een gedeeltelijk ruizig beeld $xt$ neemt — een lineaire interpolatie tussen pure ruis $x0$ en een schoon beeld $x_1$ op tijdstip $t \in [0,1]$:
$$xt = t \cdot x1 + (1 - t) \cdot x_0$$
Het model voorspelt vervolgens de richting (snelheid) naar het schone beeld:
$$v = x1 - x0$$
Tijdens de inference springt de sampler niet direct naar het doelbeeld $x_1$, maar zet hij een fractionele (kleine) stap in de richting van $v$ en herhaalt dit proces enkele keren, waardoor fijnere details kunnen ontstaan. Ik heb ook geprobeerd het model te reflowen (self-distilling) om generatie in nog minder stappen mogelijk te maken, maar besloot dit uiteindelijk niet te doen omdat de generatietijd voor 8 stappen nog acceptabel was.
Classifier Free Guidance (CFG)
Een zeer significante verbetering was de introductie van classifier free guidance (CFG). Het idee is simpel: de output van het model (snelheid) wordt één keer geëvalueerd geconditioneerd op de doelklasse en één keer met een ongeconditioneerde (null) klasse:
$$v = v{\text{null}} + w \cdot (v{\text{cond}} - v_{\text{null}})$$
In feite worden richtingen naar de doelklasse geïsoleerd en versterkt. Dit verdubbelt de verwerkingstijd, maar de resultaten zijn indrukwekkend.
Diffusion Transformer (DiT)
Aanvankelijk probeerde ik een klassiek convnet-gebaseerd U-Net, omdat ik verwachtte dat dit beter zou presteren op een klein apparaat, maar een transformer-model bleek de schonere en beter presterende keuze. De latente codering van het beeld wordt 'patchified' in $8 \times 8 = 64$ tokens van dimensie 128 door een embedding-matrix toe te passen op elke $2 \times 2 \times 8$ patch van latents. Er wordt een sinusvormige positionele embedding toegepast op elke patch (geleerde embeddings presteerden niet goed).
Een opvallende modificatie was de introductie van een $\text{ReLU}^2$ activatiefunctie in de dense blocks. Dit verhoogt de activiteits-sparsiteit, wat de inference-engine exploiteert om de inferentietijd met ongeveer 15% te verminderen.
AdaLN-Zero
Het model gebruikt AdaLN-Zero om tijdstap- en klassecondities toe te passen. Dit is een verrassend eenvoudige aanpak waarbij een geleerde bias en schaal worden toegepast op de normalisatiefuncties. Aangezien er slechts relatief weinig combinaties van tijdstappen (8) en klassen (5) nodig zijn, is het mogelijk om de getrainde AdaLN MLP weg te gooien en de biases simpelweg op te slaan in een lookup-tabel. Achteraf gezien zou AdaLN-single de tabelgrootte nog verder kunnen hebben verkleind.
Kwantisering
Het model is getraind in floating point en er is post-training kwantisering toegepast om de gewichten in int8 op te slaan. Hierbij ontstonden diverse schalingsproblemen; vooral het grotere model leed onder aanzienlijke kwantiseringsschade. Een deel hiervan kon worden hersteld door quantization-aware self-distillation na de kwantisering. Er is nog veel ruimte voor verbetering door het model op een meer kwantiseringsvriendelijke manier te trainen en de activiteitsdistributie beter te beheersen.
Eindmodel Grootte
| Sectie | m3longcfg (2.567.828 B) | m3decDdeep_full (4.016.632 B) |
|---|
| DiT block gewichten (int8) | 1.656.832 | 2.482.416 |
| Conditionering stap-tabellen (5 cond × 8 stappen × diepte) | 737.280 | 983.040 |
| VAE decoder (int8) | 117.603 | 498.411 |
| Positionele embedding (int8) | 16.384 | 16.384 |
| Final-norm gain/bias tabellen | 20.480 | 20.480 |
| Schedule, LUTs, scales, overig | ~19.000 | ~16.000 |
Het grote model past, inclusief de inference-engine, nog steeds in 4 MB flashgeheugen.
Inference Engine
De DiT is zeer geschikt voor weight streaming vanuit flashgeheugen, aangezien alle 64 tokens per laag parallel kunnen worden verwerkt. Gewichten worden via DMA vanuit flash naar een ping-pong buffer gestreamd terwijl de vorige laag nog wordt verwerkt. Omdat elk gewicht in een laag 64 keer wordt hergebruikt, wordt de inference doorgaans niet beperkt door de doorvoersnelheid van de gewichtsstreaming.
Cortex-M33 intrinsics worden gebruikt waar mogelijk (bijv. SMLAD voor parallelle vermenigvuldiging). Beide cores worden gebruikt op 300 MHz (overgeklokt vanaf de standaard 150 MHz). Iteratieve optimalisaties hebben de verwerkingstijd per beeld verminderd van de eerste implementatie naar ongeveer 5 seconden. Het grotere model en CFG brachten dit weer terug naar ongeveer 20 seconden. Desondanks is dit verbazingwekkend snel voor een diffusiemodel op een microcontroller.
Output
Het diffusiemodel wordt aangestuurd via een virtuele COM-poort op de USB-poort met een eenvoudig eigen protocol.
Daarnaast worden beelden na generatie getoond op een VGA-monitor. Er wordt gebruikgemaakt van Floyd-Steinberg dithering om banding op de 12-bit VGA-output te verminderen. Tijdens de generatie wordt een lage-resolutie preview van de latents getoond; deze wordt puur via DMA gerenderd met een kleine geheugenvoetafdruk, waardoor de inference nauwelijks wordt beïnvloed. Het SRAM dat door de framebuffer wordt gebruikt, dient tijdens de inference als scratchpad voor activaties en gewichten.
Conclusies
Dit was een opwindend project met vele nachtelijke trainingsruns die mijn kantoor hebben verwarmd. Ik ben erg tevreden met het resultaat. Er zijn nog volop mogelijkheden voor verdere optimalisatie: AdaLN-single, volledige quantization aware training en het optimaliseren van de modelarchitectuur.
***
Bronnen en referenties:
- R. Rombach, A. Blattmann, D. Lorenz, P. Esser, B. Ommer, High-Resolution Image Synthesis with Latent Diffusion Models, CVPR 2022. arXiv:2112.10752.
- R. Zhang, P. Isola, A. A. Efros, E. Shechtman, O. Wang, The Unreasonable Effectiveness of Deep Features as a Perceptual Metric, CVPR 2018. arXiv:1801.03924.
- Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, M. Le, Flow Matching for Generative Modeling, ICLR 2023. arXiv:2210.02747.
- X. Liu, C. Gong, Q. Liu, Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, ICLR 2023. arXiv:2209.03003.
- E. Heitz, L. Belcour, T. Chambon, Iterative $\alpha$-(de)Blending: a Minimalist Deterministic Diffusion Model, SIGGRAPH 2023. arXiv:2305.03486.
- Flow Matching YouTube Video.
- J. Ho, T. Salimans, Classifier-Free Diffusion Guidance, NeurIPS 2021 Workshop on Deep Generative Models. arXiv:2207.12598.
- D. R. So, W. Mańke, H. Liu, Z. Dai, N. Shazeer, Q. V. Le, Primer: Searching for Efficient Transformers for Language Modeling, NeurIPS 2021. arXiv:2109.08668.
- W. Peebles, S. Xie, Scalable Diffusion Models with Transformers, ICCV 2023. arXiv:2212.09748.
- J. Chen et al., PixArt-$\alpha$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis, ICLR 2024. arXiv:2310.02748.
AI-beeldgeneratie op een RP2350-microcontroller
Onlangs ben ik naar het project teruggekeerd en ik ben erin geslaagd een generatief AI-beeldmodel te implementeren op een RP2350-microcontroller van 1 dollar, zoals gebruikt in de Raspberry Pi Pico 2. Het model en de inference-code zijn kleiner dan 4 MB en draaien op de dual-core Cortex-M33 met 520 KB RAM.
Wat doet het?
De hardware bestaat uit een RP2350-microcontroller op een Waveshare RP2350-Plus board, aangesloten op een Pimoroni VGA-board. Het VGA-board voegt in feite een resistor-gebaseerde DAC en een VGA-connector toe, waardoor aansluiting op een VGA-monitor mogelijk is. De monitor is niet strikt noodzakelijk, aangezien het gegenereerde beeld ook via de USB-verbinding kan worden uitgelezen. HDMI-ondersteuning zou evenzeer mogelijk zijn, maar ik had geen passend board bij de hand.
Het apparaat genereert en toont RGB-afbeeldingen van menselijke gezichten met een resolutie van 128×128 pixels, waarbij elke afbeelding 10 tot 20 seconden in beslag neemt. Deze worden vervolgens weergegeven op een VGA-monitor of gestreamd via USB.
Het model implementeert een latent flow diffusion transformer (DiT), vergelijkbaar met wat wordt gebruikt in diffusiemodellen zoals Flux. Er zijn twee varianten met respectievelijk 2,9 en 1,7 miljoen parameters, wat ongeveer 4000 keer minder is dan Flux. Het ondersteunt conditionele generatie in vijf klassen: vier combinaties van geslacht en glimlach, plus een ongeconditioneerde klasse. Het model is getraind op de FFHQ-dataset.
Het is verbazingwekkend dat een model van deze omvang in staat is om complexe beelden te genereren. Veel MNIST-speelgoedprojecten voor diffusie gebruiken veel meer parameters en zijn nauwelijks in staat om iets coherents te genereren.
Hoe werkt het?
Hoewel Claude Code met Fable 5 een groot deel van het zware werk bij de implementatie van de code heeft gedaan, nam de ontwikkeling bijna twee weken in beslag, inclusief vele ablaties en nachtelijke trainingsruns op een RTX 5090. Wat me opviel, is dat de meeste optimalisaties die hielpen bij grote modellen ook noodzakelijk waren om dit micro-model werkend te krijgen.
Latente Diffusie
Het model genereert geen beelden in de pixelruimte, maar genereert een gecomprimeerde latente representatie van het beeld. Deze aanpak werd gepionierd door Stable Diffusion (Rombach et al.) en is nu in feite de standaard voor diffusiemodellen. Voor dit model wordt de pixelruimte van 128×128×3 met een factor 24 gecomprimeerd naar een 16×16×8 representatie met behulp van een variational auto-encoder (VAE).
De VAE is afzonderlijk getraind op de FFHQ-dataset. In een ongebruikelijke opstelling is de encoder veel groter dan de decoder. De reden hiervoor was om de grote encoder een latente representatie te laten leren die de decodering optimaliseert voor een kleine, on-device decoder. Voor de inference op de microcontroller is alleen de decoder nodig. LPIPS werd gebruikt als visuele perceptuele verliesfunctie.
Er zijn twee VAE's getraind: één met 115k parameters en één met 494k parameters, respectievelijk voor het 'snelle' en het 'kwaliteitsmodel'. De decoder met 494k parameters is slechts marginaal beter (bijvoorbeeld in de weergave van tanden), wat aantoont dat de latente dimensie de beperkende factor is.
Flow-based Diffusiemodel
Diffusiemodellen worden vaak verpakt in onnodig complexe wiskundige taal (DDPM, DDIM...), terwijl ze relatief eenvoudig zijn. Ik heb een flow-matching objectief gebruikt om de diffusie-transformer te trainen.
In essentie trainen we een model dat een gedeeltelijk ruizig beeld $xt$ neemt — een lineaire interpolatie tussen pure ruis $x0$ en een schoon beeld $x_1$ op tijdstip $t \in [0,1]$:
$$xt = t \cdot x1 + (1 - t) \cdot x_0$$
Het model voorspelt vervolgens de richting (snelheid) naar het schone beeld:
$$v = x1 - x0$$
Tijdens de inference springt de sampler niet direct naar het doelbeeld $x_1$, maar zet hij een fractionele (kleine) stap in de richting van $v$ en herhaalt dit proces enkele keren, waardoor fijnere details kunnen ontstaan. Ik heb ook geprobeerd het model te reflowen (self-distilling) om generatie in nog minder stappen mogelijk te maken, maar besloot dit uiteindelijk niet te doen omdat de generatietijd voor 8 stappen nog acceptabel was.
Classifier Free Guidance (CFG)
Een zeer significante verbetering was de introductie van classifier free guidance (CFG). Het idee is simpel: de output van het model (snelheid) wordt één keer geëvalueerd geconditioneerd op de doelklasse en één keer met een ongeconditioneerde (null) klasse:
$$v = v{\text{null}} + w \cdot (v{\text{cond}} - v_{\text{null}})$$
In feite worden richtingen naar de doelklasse geïsoleerd en versterkt. Dit verdubbelt de verwerkingstijd, maar de resultaten zijn indrukwekkend.
Diffusion Transformer (DiT)
Aanvankelijk probeerde ik een klassiek convnet-gebaseerd U-Net, omdat ik verwachtte dat dit beter zou presteren op een klein apparaat, maar een transformer-model bleek de schonere en beter presterende keuze. De latente codering van het beeld wordt 'patchified' in $8 \times 8 = 64$ tokens van dimensie 128 door een embedding-matrix toe te passen op elke $2 \times 2 \times 8$ patch van latents. Er wordt een sinusvormige positionele embedding toegepast op elke patch (geleerde embeddings presteerden niet goed).
Een opvallende modificatie was de introductie van een $\text{ReLU}^2$ activatiefunctie in de dense blocks. Dit verhoogt de activiteits-sparsiteit, wat de inference-engine exploiteert om de inferentietijd met ongeveer 15% te verminderen.
AdaLN-Zero
Het model gebruikt AdaLN-Zero om tijdstap- en klassecondities toe te passen. Dit is een verrassend eenvoudige aanpak waarbij een geleerde bias en schaal worden toegepast op de normalisatiefuncties. Aangezien er slechts relatief weinig combinaties van tijdstappen (8) en klassen (5) nodig zijn, is het mogelijk om de getrainde AdaLN MLP weg te gooien en de biases simpelweg op te slaan in een lookup-tabel. Achteraf gezien zou AdaLN-single de tabelgrootte nog verder kunnen hebben verkleind.
Kwantisering
Het model is getraind in floating point en er is post-training kwantisering toegepast om de gewichten in int8 op te slaan. Hierbij ontstonden diverse schalingsproblemen; vooral het grotere model leed onder aanzienlijke kwantiseringsschade. Een deel hiervan kon worden hersteld door quantization-aware self-distillation na de kwantisering. Er is nog veel ruimte voor verbetering door het model op een meer kwantiseringsvriendelijke manier te trainen en de activiteitsdistributie beter te beheersen.
Eindmodel Grootte
| Sectie | m3longcfg (2.567.828 B) | m3decDdeep_full (4.016.632 B) |
|---|
| DiT block gewichten (int8) | 1.656.832 | 2.482.416 |
| Conditionering stap-tabellen (5 cond × 8 stappen × diepte) | 737.280 | 983.040 |
| VAE decoder (int8) | 117.603 | 498.411 |
| Positionele embedding (int8) | 16.384 | 16.384 |
| Final-norm gain/bias tabellen | 20.480 | 20.480 |
| Schedule, LUTs, scales, overig | ~19.000 | ~16.000 |
Het grote model past, inclusief de inference-engine, nog steeds in 4 MB flashgeheugen.
Inference Engine
De DiT is zeer geschikt voor weight streaming vanuit flashgeheugen, aangezien alle 64 tokens per laag parallel kunnen worden verwerkt. Gewichten worden via DMA vanuit flash naar een ping-pong buffer gestreamd terwijl de vorige laag nog wordt verwerkt. Omdat elk gewicht in een laag 64 keer wordt hergebruikt, wordt de inference doorgaans niet beperkt door de doorvoersnelheid van de gewichtsstreaming.
Cortex-M33 intrinsics worden gebruikt waar mogelijk (bijv. SMLAD voor parallelle vermenigvuldiging). Beide cores worden gebruikt op 300 MHz (overgeklokt vanaf de standaard 150 MHz). Iteratieve optimalisaties hebben de verwerkingstijd per beeld verminderd van de eerste implementatie naar ongeveer 5 seconden. Het grotere model en CFG brachten dit weer terug naar ongeveer 20 seconden. Desondanks is dit verbazingwekkend snel voor een diffusiemodel op een microcontroller.
Output
Het diffusiemodel wordt aangestuurd via een virtuele COM-poort op de USB-poort met een eenvoudig eigen protocol.
Daarnaast worden beelden na generatie getoond op een VGA-monitor. Er wordt gebruikgemaakt van Floyd-Steinberg dithering om banding op de 12-bit VGA-output te verminderen. Tijdens de generatie wordt een lage-resolutie preview van de latents getoond; deze wordt puur via DMA gerenderd met een kleine geheugenvoetafdruk, waardoor de inference nauwelijks wordt beïnvloed. Het SRAM dat door de framebuffer wordt gebruikt, dient tijdens de inference als scratchpad voor activaties en gewichten.
Conclusies
Dit was een opwindend project met vele nachtelijke trainingsruns die mijn kantoor hebben verwarmd. Ik ben erg tevreden met het resultaat. Er zijn nog volop mogelijkheden voor verdere optimalisatie: AdaLN-single, volledige quantization aware training en het optimaliseren van de modelarchitectuur.
***
Bronnen en referenties:
- R. Rombach, A. Blattmann, D. Lorenz, P. Esser, B. Ommer, High-Resolution Image Synthesis with Latent Diffusion Models, CVPR 2022. arXiv:2112.10752.
- R. Zhang, P. Isola, A. A. Efros, E. Shechtman, O. Wang, The Unreasonable Effectiveness of Deep Features as a Perceptual Metric, CVPR 2018. arXiv:1801.03924.
- Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, M. Le, Flow Matching for Generative Modeling, ICLR 2023. arXiv:2210.02747.
- X. Liu, C. Gong, Q. Liu, Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, ICLR 2023. arXiv:2209.03003.
- E. Heitz, L. Belcour, T. Chambon, Iterative $\alpha$-(de)Blending: a Minimalist Deterministic Diffusion Model, SIGGRAPH 2023. arXiv:2305.03486.
- Flow Matching YouTube Video.
- J. Ho, T. Salimans, Classifier-Free Diffusion Guidance, NeurIPS 2021 Workshop on Deep Generative Models. arXiv:2207.12598.
- D. R. So, W. Mańke, H. Liu, Z. Dai, N. Shazeer, Q. V. Le, Primer: Searching for Efficient Transformers for Language Modeling, NeurIPS 2021. arXiv:2109.08668.
- W. Peebles, S. Xie, Scalable Diffusion Models with Transformers, ICCV 2023. arXiv:2212.09748.
- J. Chen et al., PixArt-$\alpha$: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis, ICLR 2024. arXiv:2310.02748.
Groetjes,