17 KiB
Ginka 地图生成模型重构设计文档
背景与动机
现状与问题
当前模型采用**热力图（Heatmap）**作为生成条件，引导 MaskGIT 模型生成地图。实践中发现以下问题：
- 模型侧：MaskGIT 无法有效利用热力图信息，热力图对生成结果的控制力弱；
- 用户侧：用户难以手动构造语义合理的热力图，交互成本高；
- 表达侧：热力图是高维连续信号，携带冗余信息，但实际指导生成的有效信息量有限。
改进目标
- 保留 MaskGIT 的地图生成范式，但将生成条件改为低维离散隐变量 z；
- z 由 VQ-VAE 风格的编码器从真实地图中抽取，形成一个可复用的 codebook；
- 训练期间通过 codebook 为 MaskGIT 提供多样性控制信号；
- 推理期间直接随机采样 z，无需用户提供任何额外输入。
整体架构
系统由两个子模型组成，训练可以分阶段进行，也可以联合训练。
┌─────────────────────────────────────────────────────────┐
│ 训练阶段 │
│ │
│ 真实地图 ──► [VQ-VAE 编码器] ──► z（离散码序列） │
│ │ │
│ ▼ │
│ [Codebook] │
│ │ │
│ ▼ │
│ 掩码地图 + z ──► [MaskGIT] ──► 预测被遮盖的图块 │
└─────────────────────────────────────────────────────────┘
┌─────────────────────────────────────────────────────────┐
│ 推理阶段 │
│ │
│ 随机采样 z ──► [MaskGIT 迭代解码] ──► 生成地图 │
└─────────────────────────────────────────────────────────┘
子模型一：VQ-VAE 风格编码器
规模约束：编码器参数量应控制在 < 1M。此前 MaskGIT 在热力图方案中已验证 4M 模型处于欠拟合与过拟合的边界，改用新方案整体规模不会有大幅变化，VQ-VAE 编码器作为辅助模块无需过大。
职责
将一张完整的地图（[H, W] 整数矩阵）编码为一个离散的 z，z 是从 codebook 中查得的嵌入向量（或向量序列）。
设计方案
编码器统一采用纯 Transformer 结构：地图尺寸（13×13）不大，Transformer 不会带来过大计算开销，同时更擅长捕捉图块之间的长程依赖，优于 CNN 方案。
方案 A：全图单一码字（z 为标量 index）
- 编码器将整张地图压缩为一个
d_z维向量，再量化为 codebook 中的单个码字； - z 是一个整数 index，对应 codebook 中一行嵌入；
- 结构最简单，但表达能力有限，难以捕捉地图的局部多样性。
地图 [H, W] ──► Transformer 编码 ──► [d_z] ──► 量化 ──► index（标量）
│
▼
codebook[index] → z [d_z]
方案 B：空间码字序列（z 为序列，已选定）
- 编码器将地图编码为
[L, d_z]的特征序列（L 为码字数量）； - 每个位置独立量化为 codebook 中的一个码字；
- z 是
L个 index 的组合，表达能力更强； - 推理时随机采样 L 个 index 即可。
地图 [H, W] ──► Transformer 编码 ──► [L, d_z] ──► 逐位量化 ──► [L 个 index]
│
▼
codebook[index_0..L] → z [L, d_z]
VQ-VAE 量化机制
标准向量量化（Vector Quantization）：
z_q = \arg\min_{e_k \in \mathcal{E}} \| z_e - e_k \|_2其中 \mathcal{E} = \{e_1, ..., e_K\} 为 codebook，K 为码本大小。
**直通估计（Straight-Through Estimator）**用于反向传播：
\frac{\partial \mathcal{L}}{\partial z_e} \approx \frac{\partial \mathcal{L}}{\partial z_q}损失函数设计
总损失由三部分组成：
\mathcal{L}_{VQVAE} = \mathcal{L}_{recon} + \beta \cdot \mathcal{L}_{commit} + \gamma \cdot \mathcal{L}_{uniform}| 损失项 | 说明 |
|---|---|
\mathcal{L}_{recon} |
重建损失，确保 z 能恢复地图信息（可选，视是否加解码器而定） |
\mathcal{L}_{commit} |
承诺损失 $|z_e - \text{sg}(z_q)|^2$，拉近编码向量与码字 |
\mathcal{L}_{uniform} |
均匀分布正则化，鼓励各码字被均等使用（见下节） |
均匀分布正则化
目标：让所有 K 个码字在训练中被均等使用，避免 codebook collapse。
方案一：熵最大化
\mathcal{L}_{uniform} = -H(p) = \sum_{k=1}^{K} p_k \log p_k其中 p_k 为码字 k 在当前 batch 中被选中的频率（使用 EMA 估计）。
方案二：EMA 更新 + 重置
- 使用指数移动平均（EMA）更新 codebook，不通过梯度更新；
- 定期检测使用率低于阈值的码字，将其重置为当前 batch 中随机样本点。
方案三：codebook 使用 KL 散度
\mathcal{L}_{uniform} = \text{KL}(p \| \mathcal{U}_{K})其中 \mathcal{U}_K 为均匀分布，即期望每个码字被选中概率为 $1/K$。
初始方案：采用熵最大化（方案一），通过梯度直接优化，实现简单；若后续出现 codebook collapse 问题，再引入 EMA 更新 + 重置机制。
是否需要解码器
| 有解码器 | 无解码器 | |
|---|---|---|
| 重建损失 | 有，z 需要还原地图 | 无，z 的质量由 MaskGIT 的生成质量间接约束 |
| 训练复杂度 | 较高 | 较低 |
| z 语义性 | 强，z 明确包含地图结构信息 | 弱，z 更像风格/多样性控制 |
| 推荐场景 | 需要 z 携带结构语义 | 仅用于多样性控制 |
已确认：不加解码器，直接与 MaskGIT 联合训练。z 的语义由 MaskGIT 端的生成损失反向传播决定。z 定位为风格与多样性控制信号，而非结构重建指导——若 z 能完整重建地图，模型会忽略其他条件直接依赖 z，反而降低可控性与泛化能力。
相应地，损失函数简化为：
\mathcal{L} = \mathcal{L}_{CE}(\text{MaskGIT}) + \beta \cdot \mathcal{L}_{commit} + \gamma \cdot \mathcal{L}_{uniform}子模型二：MaskGIT 改造
现有输入
def forward(self, map: torch.Tensor, heatmap: torch.Tensor):
# map: [B, H * W]
# heatmap: [B, C, H, W]
改造后输入
def forward(self, map: torch.Tensor, z: torch.Tensor):
# map: [B, H * W] — 掩码后的地图 token 序列
# z: [B, L, d_z] — 离散隐变量（方案B）或 [B, d_z]（方案A）
移除原有的 GinkaMaskGITCond（热力图 CNN 编码器）和 gate_encoder，用 z 替代其功能。
z 注入 MaskGIT 的方式（待选）
方式一：Cross-Attention（推荐）
Transformer 解码器天然支持 cross-attention：
map token sequence ──► self-attention ──► cross-attention ◄── z
│
▼
预测 logits
z 作为 memory 输入到 TransformerDecoder，map token 序列作为 query。
优点：z 可以是任意长度的序列（方案B友好），表达力强。已选定此方案。
方式二：z 拼接到序列头部（Prefix Token）
将 z 的每个码字嵌入作为前缀 token，拼接到 map token 序列前：
[z_0, z_1, ..., z_L, tile_0, tile_1, ..., tile_{H*W}]
输入到 TransformerEncoder 做自注意力，位置编码需相应扩展。
优点：实现简单，无需 encoder-decoder 分离。
方式三：Add/FiLM 调制
将 z 投影后，用 FiLM（Feature-wise Linear Modulation）调制每一层的特征：
y = \gamma(z) \odot x + \beta(z)优点：参数高效；缺点：z 必须先聚合为单向量（适合方案A）。
用户使用场景
调研结果显示，用户主要有以下两种使用意图：
| 场景 | 描述 | 对应训练子集 |
|---|---|---|
| 完全随机生成 | 用户不提供任何条件，直接生成完整地图 | 子集 C（随机墙壁 + Mask） |
| 墙壁辅助生成 | 用户手绘墙壁结构，模型填充非墙内容 | 子集 B / D |
训练策略需针对这两种场景显式建模，通过多子集混合训练使单一模型同时支持两种工作模式。
训练策略
多子集混合训练
每个 batch 从四个子集中按比例采样，各子集对模型能力的针对性不同：
子集 A：标准 MaskGIT 训练（比例 0.4 ~ 0.6）
保留原有 MaskGIT 掩码范式，随机遮盖部分 tile，训练模型基本的补全能力。
原始地图 ──► 随机 Mask 部分 tile ──► MaskGIT 预测被遮盖 tile
- 掩码策略：沿用现有
MapMask（随机掩码 + 分块掩码 + 形态学变换） - 作用：维持模型的基础生成质量，防止针对性训练导致退化
子集 B：墙壁条件生成（比例 0.1 ~ 0.3）
去除地图中所有非墙 tile（保留墙壁和空地），让模型在给定墙壁结构下生成其余内容。
原始地图 ──► 清除所有非墙 tile（保留 空地/墙壁）──► MaskGIT 生成非墙内容
- 输入：仅含墙壁（tile=1）和空地（tile=0）的地图
- 目标：模型输出完整的地图（含怪物、道具、门、钥匙等）
- 对应场景：用户手绘墙壁 → 模型填充
子集 C：随机条件生成（比例 0.1 ~ 0.3）
在子集 B 的基础上，进一步对墙壁进行随机 Mask，仅保留部分墙壁作为初始条件，引入随机性以支持完全随机生成场景。
原始地图 ──► 清除非墙 tile ──► 随机 Mask 部分墙壁 tile ──► MaskGIT 生成
- 输入：随机保留的稀疏墙壁片段
- 目标：生成完整地图（含墙壁补全 + 非墙内容填充）
- 对应场景：完全随机生成（墙壁本身也需要生成）
- 注：z 的随机采样在此子集中发挥关键作用，控制生成风格多样性
子集 D：入口条件生成（比例 0.1 ~ 0.2）
在子集 B / C 的基础上额外保留入口 tile（tile=10），训练模型在给定入口位置时调整生成结果。
原始地图 ──► 清除非墙非入口 tile（保留 空地/墙壁/入口）──► MaskGIT 生成
- 输入：墙壁结构 + 入口位置
- 目标：生成与入口位置语义一致的完整地图
- 对应场景：用户指定入口 → 模型生成与之匹配的布局
- 注：单独设立此子集的原因是入口对地图整体拓扑结构有较强约束，需要专项强化
各子集比例汇总
| 子集 | 训练任务 | 建议比例 |
|---|---|---|
| A | 标准 MaskGIT（随机掩码补全） | 0.4 ~ 0.6 |
| B | 墙壁条件生成（保留全部墙壁） | 0.1 ~ 0.3 |
| C | 随机条件生成（随机保留部分墙壁） | 0.1 ~ 0.3 |
| D | 入口条件生成（保留墙壁 + 入口） | 0.1 ~ 0.2 |
各子集比例之和为 1，可作为 dataset.py 中采样权重使用。
联合训练流程
将 VQ-VAE 编码器与改造后的 MaskGIT 一起训练：
真实地图 ──► VQ-VAE 编码器 ──► z（离散）
│
按子集构造的掩码地图 + z ──► MaskGIT ──► 预测 logits ──► 交叉熵损失
│
▼
VQ 损失（commit + uniform）
总损失：
\mathcal{L} = \mathcal{L}_{CE}(\text{MaskGIT}) + \beta \cdot \mathcal{L}_{commit} + \gamma \cdot \mathcal{L}_{uniform}Dropout z（提升鲁棒性）
训练时以一定概率（如 10-20%）将 z 替换为随机采样的码字，模拟推理时随机采样 z 的情况，避免模型过度依赖 z 的精确语义。在子集 C（完全随机生成）中此机制尤为重要。
推理流程
场景一：完全随机生成
1. 随机从 [0, K) 均匀采样 L 个 index
2. 从 codebook 查表，得到 z [L, d_z]
3. 初始化地图序列：以少量随机墙壁（比例 0.02 ~ 0.1）作为初始条件，其余位置填充 MASK token
4. 迭代 MaskGIT 解码（cosine schedule）：
a. 输入掩码地图 + z → 预测所有 MASK 位置的 logits
b. 取置信度最高的 N 个 token 进行 unmask
c. 重复直到无 MASK token
5. 输出最终地图
初始墙壁比例说明：全 MASK 初始化会导致生成多样性不佳（即使使用 top-k 采样也难以改善）。引入少量随机墙壁（2%~10%）作为种子，可有效提升多样性；但比例不宜过高，否则随机放置的墙壁可能产生拓扑结构冲突，影响生成质量。
场景二：墙壁辅助生成（用户手绘墙壁）
1. 用户提供墙壁布局（tile=1 的位置），其余位置填充 MASK token
2. 随机从 [0, K) 均匀采样 L 个 index，得到 z
3. 迭代 MaskGIT 解码：
a. 输入带已知墙壁的掩码地图 + z → 预测 MASK 位置的 logits
b. 已知墙壁位置不参与 unmask，保持不变
c. 取置信度最高的 N 个 MASK token 进行 unmask
d. 重复直到无 MASK token
4. 输出最终地图
场景三：入口指定生成
与场景二类似，但初始条件同时固定墙壁位置和入口位置（tile=10）
其余 MASK 位置由 MaskGIT 迭代解码填充
超参数（待确定）
| 参数 | 说明 | 建议初始值 |
|---|---|---|
K |
codebook 大小 | 8 ~ 32 |
L |
码字序列长度（方案B） | 1 ~ 4 |
d_z |
码字嵌入维度 | 64 ~ 128 |
β |
commit loss 权重 | 0.25 |
γ |
uniform loss 权重 | 0.1 |
| z dropout 概率 | 训练时随机 z 的比例 | 0.1 ~ 0.2 |
地图固定为 13×13，z 定位为风格与多样性控制信号，K 无需过大（8~32 足够），L 也应保持较小以避免过度约束生成。
已决定事项
| 决策点 | 结论 |
|---|---|
| VQ-VAE 解码器 | 不加，直接与 MaskGIT 联合训练，z 语义由生成损失反向传播决定 |
| 编码器架构方案 | 方案 B（序列码字），全 Transformer 结构；若训练不稳定可退回方案 A |
| z 注入方式 | Cross-Attention（方式一） |
| 均匀正则化策略 | 熵最大化；若出现 collapse 再引入 EMA |
| 分层 VQ | 单层，13×13 地图无需多级量化 |
| 标量 cond 保留 | 暂不使用；训练集标量分布严格，推理阶段难以生成合理值，后续可再考虑加入 |
待探索事项
- 合适的 K、L 取值（建议从 K=16, L=2 开始实验），K=1, L=64 可能比较合适。
- z dropout 的最优概率
- 若后续 codebook collapse：引入 EMA 更新 + 重置机制
- 若后续需要更细粒度控制：加入标量 cond（需对推理侧标量做随机扰动处理）
下一步行动
- 确定编码方案（方案 B，序列码字）
- 确定 z 注入方式（Cross-Attention）
- 确定是否需要 VQ-VAE 解码器（不加）
- 确定是否保留标量 cond（暂不使用）
- 确定训练子集划分与各场景比例
- 实现 VQ-VAE 编码器模块（Transformer + VQ）
- 改造
GinkaMaskGIT.forward()接受 z，替换热力图分支 - 实现四种子集采样逻辑（
dataset.py新增多子集 Dataset 或采样权重） - 实现子集 B/C/D 的输入构造函数（按规则清除 tile、保留墙壁/入口）
- 编写联合训练脚本，整合 VQ 损失与 MaskGIT 交叉熵损失