极致压缩：170 万参数扩散模型在 1 美元单片机上跑通

概述

独立开发者将一套完整的潜在扩散 Transformer（DiT）生图模型，移植到售价约 1 美元的 RP2350 单片机上。这块芯片正是树莓派 Pico 2 使用的型号，运行在双核 Cortex-M33 与 520 KB RAM 之上。整套模型与推理代码不到 4 MB，能在 10–20 秒内生成一张 128×128 RGB 人脸图像，并通过 VGA 显示器或 USB 输出。这一成果展示了将生成式 AI 压缩到极端受限设备中的可行路径。

模型与功能

该项目的核心是一个潜在空间流匹配扩散 Transformer，工作机制与 Flux 等扩散模型类似，但参数规模大幅压缩：

• 两个变体：快速版 170 万参数、质量版 290 万参数；
• 与 Flux 相比，参数规模约为后者的 1/4000；
• 支持 5 个条件类别的生成：男性微笑、女性微笑、男性不笑、女性不笑，以及一个无条件类别；
• 训练数据采用 FFHQ 人脸数据集。

在硬件层面，开发者使用了 Waveshare 的 RP2350-Plus 开发板，配合 Pimoroni 的 VGA 子板（自带电阻式 DAC 与 VGA 接口），因此不依赖显示器时也可通过 USB 读取生成结果。

关键技术路径

潜在空间扩散

模型并不直接在像素空间生图，而是先编码到潜在空间后再做扩散，这一点沿用了 Stable Diffusion 的思路。具体做法是将 128×128×3 的像素空间通过 VAE 以 24 倍压缩率映射为 16×16×8 的潜在表示。VAE 与扩散模型分开训练：编码器较大（用于学习易于解码的潜在表示），解码器较小（用于设备端推理）。两个 VAE 分别训练了 11.5 万参数与 49.4 万参数两个版本，配合快速版与质量版模型使用，并以 LPIPS 作为感知损失。

流匹配目标

扩散部分采用流匹配（flow-matching）目标训练 DiT：模型接收带噪的插值图像 x_t，预测指向干净图像的速度场 v = x_1 - x_0，推理时按小步长迭代更新。开发者提到该目标比传统 DDPM/DDIM 表述更简洁，并尝试过 reflow（自蒸馏）来减少步数，但最终 8 步生成耗时已可接受。

无分类器引导（CFG）

引入 Classifier Free Guidance 后效果提升明显：模型同时以条件类与无条件类各推理一次，将条件方向的差量放大叠加。尽管这会使推理时间翻倍，但对生成质量的改善相当显著。

DiT 架构与小优化

出乎意料的是，在这种微型规模下 Transformer 比传统 U-Net 表现更好。潜在表示被切成 8×8=64 个 token，并使用正弦位置编码（学习型位置编码效果不佳）。关键改动之一是在全连接层使用 ReLU² 激活，推理引擎可借此利用激活稀疏性，将推理时间缩短约 15%。

开发过程与硬件约束

• 推理目标硬件为双核 Cortex-M33 + 520 KB RAM，开发板为 RP2350-Plus；
• VAE 编码器明显大于解码器，仅解码器进入设备；
• 模型构建与调试借助 Claude Code（搭配 Fable 5）完成主要代码工作；
• 消融实验与每晚训练在 RTX 5090 上进行，整体耗时约两周。

这套实现证明，原本为大模型设计的扩散优化经验，同样适用于参数量极少的微型模型。同时也意味着，在未来更多低成本 MCU 上跑动小型生成式模型具备现实可行性。