Aquarius — Native Ternary Text-to-Image Diffusion Model
原生三值文生图扩散模型 · 从第 0 步起，权重就活在量化空间中（Native Low-bit Training, NLT）
这是一个"可行性"项目，不是"可用模型"项目。我们证明了 原生低比特训练 与 低比特融合算子推理 在真实硬件上端到端可行；同时如实报告：在当前数据规模下，模型质量未达到可用发布水平，训练因指标饱和 + 经费不足在 step 260,000 处主动停止。详见下方"结果与诚实声明"。
TL;DR
| 项目 | 值 |
|---|---|
| 参数量 | 558,347,012（量化部分 540,147,712 个三值权重） |
| 量化 | 原生三值 {−s, 0, +s}，g128 mean-scale + STE（NLT，非 PTQ/QAT） |
| 交付体积 | 152.99 MB（2.190 bpw，base-3 位打包，log₂3 ≈ 1.585 bit 信息下限） |
| 推理 | 三值逐位映射进 torch int4 融合内核（q₄ = q·7 + 8 ∈ {1,8,15}），287.3 MB 常驻 |
| 训练 | COCO train2017 · 260,000 步 ≈ 28.70 epochs · 单卡 32 GB 消费级 GPU |
| 出图 | 512×512 / 20 步 DDIM，3.1~3.5 s/张，进程峰值显存 1.34 GB（RTX 5090） |
| 质量 | clip_ratio 0.2908 @ step 260k（峰值 0.3782 @ step 170k）——未达可用水平 |
这是什么？为什么要这么做？
低比特扩散模型的既有工作几乎都是"先在 fp16 预训练、再量化压缩"（PTQ / QAT）。这条路有一个结构性隐患：量化误差是在一个已经收敛到浮点解的空间里被强行引入的，模型没有机会在训练过程中适应它。
Aquarius 反过来问：如果从随机初始化的那一刻起，权重就只能取 {−s, 0, +s} 三个值，模型能学到什么？
答案（三条实测证据，非推测）：
- 原生低比特训练能跑 —— 260,000 步稳定推进，无塌缩、无发散、无 OOM；权重始终保持在三值结构内（打包校验：码位精确、无退化）。
- 低比特融合算子推理能跑 —— 三值与 int4 的逐位精确映射（q₄ = q·7 + 8），权重 287.3 MB 常驻，解包融合发生在 GEMM 内部。
- 质量不达标，且当前规模已饱和 —— clip_ratio 在 step 110k 之后 13 万步无净进展（近 8 阶段斜率 −0.00357/万步，t = −0.71）；最终读数 0.2908 未突破 step 170k 的峰值 0.3782。
模型架构
- UNet：SD1.5 同构结构（base=256，block_out_channels=[256,512,1024,1024]，cross_attention_dim=1024，8 个注意力头），仅 UNet 参与训练。
- 文本编码器：冻结的 Qwen3.5-0.8B（32 token × 1024 维），不参与训练。
- VAE：冻结，全程不参与训练。
- 采样：确定性 DDIM（η=0）+ 余弦噪声调度，与训练期一致；训练未启用 CFG，推理端无负向提示词与引导强度自由度。
量化与打包（关键设计）
- 量化器（训练期 = 推理期同一个函数）：每 128 个权重共享一个 fp16 尺度 s = mean(|w|)，q = round(clamp(w/s, −1, 1)) ∈ {−1, 0, +1}；反向用 STE。有效存储 = log₂3 + 16/128 = 1.710 bpw（该参数化下的信息下限）。
- 交付件（base-3 位打包）：三值只有 3 个状态，理论下限 log₂3 ≈ 1.585 bit；按 2 bit/权重存储浪费 26% 码位。3⁵ = 243 ≤ 256，5 个三值码无损塞进 1 个字节 → 152.99 MB / 2.190 bpw 单文件。
- 推理（int4 融合）：torch 的 int4 约定为"半字节 0..15 代表 −8..7"。令 q₄ = q·7 + 8 ∈ {1, 8, 15}，则 (q₄ − 8)·s/7 = q·s 逐位精确 → 直接调用
torch._weight_int4pack_mm，解包在 GEMM 内部，常驻 287.3 MB。 - 必避之坑：绝不能从"解包后的浮点权重"反推码位（会静默缩小到 71%，实测出图 L1 冲到 37.8/255）。正确做法是只从打包文件读 (码位, 尺度)。
结果
- 三模式对照（step 40k 同配方）：三值 clip_ratio 0.2300（152.99 MB）> 二值 0.1223（112.36 MB）> 全精度 −0.0077（1116.73 MB）。三值反而优于全精度参照——量化不是瓶颈。
- 效率：从打包文件直建模型，建模型显存峰值 1158 → 362 MB（−68.7%）；CUDA Graph 使 256² 单图延迟 1013 → 218 ms（4.65×）。
- 端到端：
diffusion_model_step260000.safetensors= 323.48 MB（运行时 int4 布局，载入零转换），512×512 / 20 步 DDIM，3.1~3.5 s/张，峰值显存 1.34 GB。
诚实声明（重要）
- 模型未达到可用发布水平：修正采样协议后，4 条 caption 中仅 1 条能稳定认出主体（如 giraffe / kitchen 场景）；结构性学习已发生，但统计效力弱（4 张 × 1 种子）。
- 训练在 step 260,000 = 28.70 轮主动停止，依据：① clip_ratio 饱和（13 万步无净进展）；② 租用算力经费不足。继续堆步数已无可测收益。
- 跨平台：
torch._weight_int4pack_mm仅注册 CUDA；AMD / Intel / Apple 设备只能退回浮点（解包 fp32 2233 MB）。移动端可行性的现实路径是 ONNX/MNN/ncnn + 自研反量化融合着色器。 - 本项目的复用价值在 §5.6 六条可迁移的实测结论（低位宽不融合是负收益、cos 不是可行性判据、恢复训练后必须核对实际生效超参等）——换一个低比特扩散项目仍然成立。
仓库内容
| 文件 | 说明 |
|---|---|
README.md |
本文件 |
| 技术报告 | 中英双版技术报告（见 attachments / 发布页） |
| 权重/源码 | 待上传（整理中） |
引用
@techreport{du2026aquarius,
title = {Aquarius: Native Ternary Text-to-Image Diffusion Model},
author = {Junyi Du},
year = {2026},
note = {Independent technical report}
}
License
本项目采用 Apache License 2.0。训练数据为 COCO train2017（Microsoft 许可条款下使用，请自行确认你的使用场景符合其条款）。