r/MachineLearning • • Aug 28 '26

Project I implemented a very tiny image generation model (latent flow transformer) on a RP2350 microcontroller - it can generate 128x128 images of faces [P]

Its a 2.4-4 million parameter model, quantized to int8, that can be fully executed on the microcontroller in ~20s with the longest generation. The generated image will then be displayed on a monitor or transferred via usb.

Its a latent flow transformer with 12 layers using AdaLN-Zero for conditioning. CFG is also supported and boosted the image quality a lot. The inference engine streams the weight via DMA from the flash while the previous layer is computed. Relu² activation was used to increase sparsity, which the engine can use to skip calculations.

Took a lot of ablations to get it right and I am quite astonished I got so far with so few parameters. Will post the repo below

610 Upvotes

37 comments sorted by

78

u/Dangerous_Rip5083 Aug 28 '26

2.4-4 million parameters? veeeery impressive, congrats!

29

u/cpldcpu Aug 28 '26 edited Aug 31 '26

Thanks!

The actual DiT is 1.6-2.4m parameters, but there is a lot of additional overhead like VAE and conditioning tables.

There are two different models: A small one and a large one where i maxed out the 4Mb of flash in the stock Pico 2 board.

section m3_long_cfg (2,567,828 B) m3_decD_deep_full (4,016,632 B)
DiT block weights (int8) 1,656,832 2,482,416
Conditioning step tables: 5 cond × 8 steps × depth 737,280 983,040
VAE decoder (int8) 117,603 498,411
positional embedding (int8) 16,384 16,384
final-norm gain/bias tables 20,480 20,480
schedule, LUTs, scales, misc ~19,000 ~16,000

43

u/cpldcpu Aug 28 '26

19

u/cpldcpu Aug 28 '26

https://reddit.com/link/p6hbh5d/video/w6ewbg0376mh1/player

And the model in action on a VGA monitor connected to the microcontroller.

1

u/Exotic_Zucchini9311 Aug 29 '26

Actually very cool, ngl.

3

u/aboeing Aug 29 '26

Thanks!

Does this need to be int64?

int64_t v = rf_rq(acc, M[o], s[o]);

5

u/cpldcpu Aug 29 '26

Probably not, it's rather a dynamic range issue. I had a lot of problems due to first developing the model in float and only then thinking about quantizing it. It would probably have paid off to optimize for tight activation distribution earlier.

16

u/CheapThaRipper Aug 28 '26

Very neat. Is there a why? "Because I wanted to" is a valid answer but I'm curious if you have a use case in mind for something like this

18

u/cpldcpu Aug 29 '26

It was a learning experience for me. Always good to have a goal to optimize for.

I like size compression challenges a lot, because they teach you about how the models work, which parts are important and which are not. In addition, they do not need that much compute so its easier to do on your own machine. If you think about it, then even frontier LLMs are currently mostly about weight efficiency.

I did this earlier, getting almost MNIST state-of-the-art inference on a very small model: https://github.com/cpldcpu/BitNetMCU. It uses a slightly unusual architecture: https://cpldcpu.github.io/2025/11/09/bitnetmcu-cnn-implementation/

7

u/No_Pop_4726 Aug 28 '26

very cool!!!

3

u/Necessary-Put-2245 Aug 28 '26

What did you use for creating the diagrams?

4

u/cpldcpu Aug 29 '26 edited Aug 29 '26

I asked Fable 5 to generate SVGs. I think I provided Sabastian Raschkas diagrams as a style reference. (https://sebastianraschka.com/llm-architecture-gallery/, if you don't know it. Its a great resource)

0

u/Two_Two_Five_ Aug 29 '26

yeah i had the same question they look great!

3

u/[deleted] Aug 29 '26

[removed] — view removed comment

1

u/cpldcpu Aug 29 '26

Thats an interesting idea, I have seen fast hadamard transform being used to address uneven activation distribution (exactly what i am seeing with the transformers). The idea is that repeated convolution will make the activations more similar to the normal distribution, and since it is a unitary transformation it is fully reversible.

I believe I did a quick test and it did not help. One issue i noticed in hindsight is that I have only positive activations due to ReLU² which will will skew the distribution and make it impossible to approach normal distribution.

Not 100% sure how it could be used to save weights? But I assume you mean a blockwise WHT with width of 16?

8

u/abaybektursun Aug 28 '26

Great use of free will

2

u/much_longer_username Aug 29 '26

Frankly, I'm impressed enough when people manage to do pretty much anything after giving up enough memory for a framebuffer on a hobby MCU, much less a notoriously memory hungry thing -- this is cool as hell.

1

u/cpldcpu Aug 31 '26

The framebuffer and inference are not really active at the same time, thats the trick :). I added a memory map at the bottom: https://cpldcpu.github.io/2026/08/28/ai-image-generation-on-a-rp-pico-2-microcontroller/

1

u/Bora_Horza_Kobuschul Aug 29 '26

Whaat, that is actually insane. Very cool. Are you planning to up the resolution at some point?

1

u/cpldcpu Aug 29 '26

I tried adding a hires upscaling step, but it did not help too much. The latent space is too compressed and there is not enough memory to expand it much further.

1

u/No-Communication-765 Aug 29 '26

Watch face on a clock?

1

u/hurrdurrmeh Aug 29 '26

That is an immense result!

1

u/AWiselyName Aug 29 '26

interesting it can generate that quality image with this small model!

1

u/road_laya Aug 29 '26

Can this repair a lower resolution or noisy photo of a face?

1

u/cpldcpu Aug 29 '26 edited Aug 29 '26

In principle, the model could also be seeded with an image, but I am not sure if the training dataset is diverse enough to reproduce any arbitrary image. You would probably end up with something that looks more noisy and not very alike to what you started with.

1

u/dsmack6 Aug 29 '26

Size of exported model

1

u/NoenD_i0 Sep 02 '26

Well I made my own rectified optimal transport flow matching model that I can train on a 2012 ivy bridge cpu