63 lines
2.2 KiB
Markdown
63 lines
2.2 KiB
Markdown
---
|
||
title: "Shared-Weight Discretization"
|
||
created: 2026-05-13
|
||
updated: 2026-05-13
|
||
type: concept
|
||
tags: [network-architecture, diffusion-language-model, multi-task-learning]
|
||
sources:
|
||
- https://arxiv.org/abs/2605.10938
|
||
---
|
||
|
||
# Shared-Weight Discretization
|
||
|
||
共享权重离散化是 [[embedded-language-flows|ELF]] 的核心设计：**同一个网络既执行去噪又执行解码，区别仅在于输入条件和输出处理**。
|
||
|
||
## 机制
|
||
|
||
网络签名为 `net_θ(z, t, mode)`，其中：
|
||
- `z`：当前（带噪）嵌入
|
||
- `t`：时间步 ∈ [0,1]
|
||
- `mode`：二进制 token，`denoise` 或 `decode`
|
||
|
||
### Denoise Mode (t < 1)
|
||
|
||
```
|
||
x̂ = net_θ(z_t, t, "denoise")
|
||
v̂ = (x̂ - z_t) / (1-t) # 转换 x-prediction 为速度
|
||
L = MSE(v̂, v_true)
|
||
```
|
||
|
||
### Decode Mode (t = 1)
|
||
|
||
```
|
||
# 先对 z 加 token 级 corruption 构造非平凡输入
|
||
z̃ = corrupt(z_1)
|
||
x̂ = net_θ(z̃, t=1, "decode")
|
||
logits = W · x̂ # unembedding 层
|
||
L = CrossEntropy(logits, s) # s 是真实 token
|
||
```
|
||
|
||
## 为什么共享权重有效
|
||
|
||
[[x-prediction-parameterization]] 是关键：网络始终预测**干净嵌入 x̂**。在 denoise mode 中它转换为速度；在 decode mode 中它直接经 unembedding 转为 logits。两种模式共享网络权重，因为它们在语义上一致——都试图恢复干净的 token 表示。
|
||
|
||
**v-prediction 无法做到这一点**：预测速度 v 与预测离散 token 之间的语义鸿沟使得权重共享不可行（ELF 论文中实验证实）。
|
||
|
||
## 优势
|
||
|
||
1. **零额外参数**：不需要单独训练的 decoder（与 LD4LG 等潜在扩散方法对比）
|
||
2. **训练效率**：两种模式在一个 batch 中通过 masking 同时训练，无额外计算开销
|
||
3. **语义对齐**：去噪目标（恢复干净嵌入）和解码目标（恢复干净 token）共享底层表示
|
||
|
||
## 实现细节
|
||
|
||
训练时两分支按比例混合（ELF 默认 80% denoise + 20% decode）。推理时：
|
||
1. t < 1：使用 denoise mode，迭代更新嵌入
|
||
2. t = 1：使用 decode mode，argmax 输出离散 token
|
||
|
||
## 相关概念
|
||
|
||
- [[embedded-language-flows]] — 使用此机制的模型
|
||
- [[x-prediction-parameterization]] — 共享权重可行的关键
|
||
- [[flow-matching]] — 使 t=1 步骤自然成为解码点的基础框架
|