Showing Posts From

Google deepmind

Brody
2026/08/05
- Tech Insights

PLE 逐层嵌入：$8 芯片是如何跑起大模型的？

如果你对 AI 模型略有关注，2026 年上半年有一个趋势已经很明显了：模型正在从云端走向端侧。无论是手机上的 Gemma 3n/E2B，还是 slvDev 在 ESP32-S3 上跑起来的 28.9M 参数 LLM，大家都在做同一件事——让小模型拥有大智慧。而其中一个关键的技术突破，叫 Per-Layer Embeddings（PLE，逐层嵌入）。这篇文章，我就来把它拆开讲清楚。标准 Transformer 有一个隐藏的瓶颈在经典的 Decoder-only Transformer 里，token 的处理方式是教科书级的：输入层做一次 embedding 查表，得到一个向量 e ∈ Rd_model 这个向量通过残差流（residual stream）贯穿所有的 Transformer 层每一层在此基础上做 Attention + FFN，逐步提取语义看起来没问题？问题出在第三步。随着层数加深，原始 token 的身份信息——「这个位置是"猫"还是"狗"」——会逐渐被后续计算稀释、混杂，甚至丢失。到了第 20 层，模型已经不太记得第 0 层输入的 token 到底是谁了。这只靠一层 embedding 表来解决，就像一个公司的前台要在你进门的那一秒把你接下来所有会议的需求都交代清楚——不现实。唯一的补偿方式是把 embedding 表做得特别大、信息密度特别高。但这也恰好是瓶颈所在：embedding 表成为小模型的主要内存杀手。vocab_size × d_model 这么大的矩阵，对于端侧芯片来说是要命的。PLE 的核心思想：给每一层一张「专属提示卡」 PLE 的思路很直觉——既然一层 embedding 不够用，那就给每一层都准备一个专属的 token 信号通道。可以这样类比：标准模型：一本书只在封面写了目录，后面章节只能靠记忆回忆 PLE 模型：每一章开头都有一张针对本章内容的「专属提示卡」，而且这张卡是自适应的具体来说，PLE 让模型在每一层都能「重新接收」该 token 的身份信息，而不是依赖输入层一次性塞入所有信息。这意味着：输入 embedding 表不必再承担「记住一切」的压力身份信息被分散到所有层，每一层只携带当前层需要的那一小部分深层计算再也「忘不掉」token 是谁技术实现：两个组件 + 一个门控 PLE 信号由两个互补部分组成，在进入 Transformer 层循环之前一次性准备好。第一步：准备 PLE 信号（1）Token-Identity 组件（纯身份信息）使用一个巨大的查找表 embed_tokens_per_layer，形状大约为： (vocab_size, num_layers × ple_dim)对当前 token ID 做一次查表，然后 reshape 成： (batch, seq_len, num_layers, ple_dim)通常还会乘以缩放因子 √ple_dim。（2）Context-Aware 组件（上下文感知信息）将标准输入 embedding 通过一个线性投影层 per_layer_model_projection 映射到相同的 packed 空间。再进行缩放（通常 1/√d_model）和 RMSNorm。（3）融合 ple_input = (token_identity + context_aware) × 1/√2得到每个层专属的切片。第二步：门控注入（最关键的一步）在每一层内部：正常执行 Attention + FFN 用当前层的 hidden state 生成一个 gate——通常是 GELU(Linear(h)) 将 gate 与该层的 PLE 向量做逐元素乘法（自适应控制注入强度）将结果投影回 d_model 维度，归一化后作为额外的 residual 加回去这意味着什么？每一层都能根据当前计算状态，有选择地吸收该 token 的身份信息。不是被动接收，而是主动调节——当前层需要多少身份信号，就拿多少。为什么这对边缘设备特别重要这才是 PLE 真正的巧思所在：它完美适配了边缘芯片的内存分层架构。以 ESP32-S3 为例，它的存储结构是这样的：存储类型速度容量用途SRAM 极快 ~512KB 频繁访问的激活值和核心权重PSRAM 中速 8MB 输出头等Flash 大但慢 16MB 巨大的 PLE 查找表PLE 表虽然大（可达 25M 参数），但它的读取模式是稀疏的——每个 token 只需要读取大约 6 行，约 450 字节。这意味着：25M 参数的大表可以安心放在 Flash 里，几乎不占用快速内存真正做矩阵乘法的「思考核心」（dense core）可以缩小到几十万到一两百万参数总参数量可以做到 28.9M，而实际活跃计算参数极小在 $8 的芯片上跑本地 LLM 成为可能这正是 slvDev/esp32-ai 项目能成功的关键原理：把大部分参数变成「只读查找表」，而不是「需要计算的权重」。优缺点一览优点显著提升小模型的有效容量：Gemma E2B 标称 2.3B 有效参数，但总参数含 embedding 达 5.1B 缓解深层 token 身份信息丢失问题，尤其是在深层仍能保持 token 级别的精度完美适配内存分层设备：手机、MCU、边缘芯片天然受益推理额外开销可控：主要是查表 + 小规模投影和门控，计算量不大缺点总参数量明显增加：大量参数存在于 PLE 表中，模型文件更大实现复杂度较高：双组件融合 + 门控注入，不是 plug-and-play 对训练数据质量和超参更敏感：门控机制需要充分训练才能稳定在超大模型上收益可能递减：大模型本身容量已足够，PLE 的边际收益下降与其他技术的关系 PLE 不是孤立的创新，它和整个端侧模型优化生态是互补的：与 ALBERT 的 factorized embeddings 有概念相似之处，但 PLE 更激进——直接给每一层独立信号与 KV-sharing、MoE 等技术是互补的：Gemma 4 同时使用了多种技术，PLE 是其中关键一环在极端边缘场景中，PLE 的「大表慢存、小核快算」理念，比单纯量化更关键量化是把模型「压缩」到能放进内存。PLE 是从架构层面改变了参数的组织方式——让大部分参数不必参与计算，而变成了可查表的知识库。一句话总结 PLE 把「token 身份知识」从单一的输入瓶颈，分散成贯穿全网络深度的低成本、可自适应的条件信号，从而让模型在极小的活跃计算预算下，仍能拥有远超其规模的知识容量。这正是它能同时服务于手机端 Gemma 和 $8 ESP32 的根本原因。参考来源：Google DeepMind Gemma 3n / Gemma 4 技术报告、slvDev/esp32-ai GitHub 仓库、Gemma 模型卡（Hugging Face）