Hotdry.
归档
ai-systems

工程化最小Transformer：精确10位数加法

基于AdderBoard挑战，剖析36参数手写与311参数训练模型的关键架构、训练参数与边缘部署清单，实现10位数精确加法。

在边缘设备上实现高效算术推理是 AI 系统工程化的关键挑战，尤其是当模型需处理如 10 位数加法这类精确任务时。传统 Transformer 因参数庞大而不适合，但 AdderBoard 项目展示了通过极致最小化架构，能以数十参数实现 100% 准确率。本文聚焦单一技术点：构建最小自回归 Transformer，支持从右到左逐位生成结果，顺带进位传播，适用于边缘算术推理。

最小架构设计：核心组件与参数清单

观点：最小 Transformer 需平衡自注意力（对齐数字）、MLP（逐位算术）和自回归生成（进位传播），通过低秩分解、权重绑定和自定义位置编码压至百参数级。

证据：AdderBoard 排行榜显示，手写权重类别冠军为 36 参数 2 层解码器（d_model=5，5 头 + 1KV 头），使用 ALiBi 位置编码（斜率 log (10) 模拟基 10 权重）和稀疏嵌入，实现 100% 准确。“ALiBi slope=log (10) for base-10 weighting” 是关键 trick，确保位置敏感性而无额外参数。

可落地参数 / 清单：

  • 嵌入层：词汇表 12（0-9 + + = PAD BOS EOS），d_model=27。技巧：因子化嵌入（rank-1/3），绑定输出头节省20% 参数。
  • 自注意力：12 层，15 头，head_dim=2。低秩 QKV（rank-3 分解），共享 A 矩阵绑定 KV。位置编码：ALiBi (log (10)) 或 RoPE (周期 19/11)，无学习 PE。
  • FFN：dim_ff=2~14，Gated ReLU/SiLU。稀疏门控、低秩投影（rank-1）。
  • 归一化：Pre-LN 或 RMSNorm，共享向量无偏置。
  • 总参数目标：<500，float64 手写或 float16 训练。
  • 序列格式：固定 35 token，零填充 10 位 +“+”=，输出反转 11 位 + EOS（进位位）。

示例 PyTorch 骨架（基于排行榜 top 模型）：

import torch.nn as nn
class TinyAdder(nn.Module):
    def __init__(self, d_model=5, n_head=5, head_dim=2, ff_dim=8):
        super().__init__()
        self.embed = nn.Embedding(12, d_model)  # 因子化可选
        # 注意力：低秩，ALiBi
        self.attn = MultiHeadAttention(d_model, n_head, head_dim, alibi=True)
        self.ffn = GatedReLU(d_model, ff_dim)
        self.norm1 = nn.RMSNorm(d_model)
        self.norm2 = nn.RMSNorm(d_model)
        self.out = nn.Linear(d_model, 12, bias=False)  # 绑定embed
        self.embed.weight = self.out.weight  # 绑定
    def forward(self, x):
        h = self.norm1(self.embed(x))
        h = h + self.attn(h)
        h = h + self.ffn(self.norm2(h))
        return self.out(h)

推理：标准自回归，温度 0，top-k=1，max_len=11。

训练效率优化：Grokking 与课程学习

观点：训练需诱导 “grokking” 相变，从过拟合到泛化，利用课程学习加速边缘 fine-tune。

证据：训练冠军 311 参数 1 层模型（d=4，rank-3），经 30K 步 grokking 达 99.999%。社区观察～800 参数悬崖：低于阈值准确率近零，高于即锐升。

可落地参数 / 清单：

  • 数据集：10^6 随机 10 位对（[0,10^10-1]），80/20 train/val。固定种子 2025 验证 10K held-out。
  • 优化：AdamW lr=1e-36e-4，warmup 10%，cosine 衰减。batch=2561024，grad_clip=1.0。
  • 课程：从小位数（1~10）渐增，避免早期卡住。损失：仅输出位 CE，忽略输入。
  • 硬件 / 时间：笔记本 GPU/MPS，5~10min 收敛。监控：train/val loss，准确率 > 99% 停止。
  • Grokking 诱导：过拟合后继续训（~4K 步零准→22K 步 99%），weight decay=0.1。

部署回滚：若 < 99%，fallback 传统加法器；监控生成 latency<1ms / 位。

边缘部署：量化、监控与扩展

观点：tiny 模型天然支持 int8 量化，结合 ONNX/TVM 至 uC 部署，实现 μW 级算术。

参数 / 清单：

  • 量化：PTQ int8（calib 1K 样本），准损 <0.1%。动态范围：embed/attn [-5,5]。
  • 推理 opt：TorchScript/jit，融合 LN+Linear。ALiBi 预计算。
  • 监控点：准确率（10K test）、latency（ms / 加法）、内存 (<10KB)、功耗 (μJ/op)。
  • 扩展：多位乘法改 FFN，泛化需 d_model>16 + 多层。
  • 风险限：仅加法，carry chain 长 > 10 易溢；手写模型非鲁棒。

此工程路径已在 AdderBoard 验证，36p 手写证明可表示性，311p 训练示 learnability。边缘 AI 从此起步。

资料来源：

查看归档