# Stage1 二维空间感知改进设计文档

## 问题诊断

### 核心现象

第一阶段（floor/wall 骨架生成）质量显著劣于第二、三阶段。生成结果常见表现为：

- 墙壁分布破碎，缺乏连通性；
- 房间边界不完整，出现孤立墙块；
- 走廊未能形成闭合通道；
- 生成结果的空间拓扑结构与训练集分布偏差较大。

第二、三阶段（门/怪物/资源放置）问题相对较少，因为这些阶段的任务是"在已有结构上填充稀疏元素"，对空间结构的整体一致性要求较低。

### 根本原因：一维位置编码与二维网格的结构失配

#### 现有架构

`GinkaMaskGIT` 当前使用以下位置编码：

```python
self.pos_embedding = nn.Parameter(torch.randn(1, map_size, d_model) * 0.02)
```

这是一个纯一维可学习位置嵌入，将 13×13=169 个格子视为一条线性序列。`GinkaVQVAE` 中同样如此。

`maskGIT.py` 中的 `Transformer` 使用标准 `nn.TransformerEncoder` + `nn.TransformerDecoder`，注意力机制中没有任何二维空间偏置。

#### 一维展平带来的结构失配

将 13×13 地图按行展平后，位置关系如下：

```
位置(0,12) → token 12
位置(1, 0) → token 13
```

这两个 token 在一维序列中相邻（距离 1），但在二维地图上相距 12 列（横跨整行）。而真正的二维邻居关系：

```
位置(0, 0) 和 位置(1, 0) 的二维距离 = 1 格
对应 token 0 和 token 13，一维距离 = 13
```

一维位置嵌入告诉模型"token 0 和 token 13 相距较远"，但实际上它们是相邻的竖向邻格。注意力机制的相对偏置完全依赖位置嵌入的初始化，无法从一维嵌入中自动推断二维邻接关系。

#### 为何 Stage1 特别敏感

Stage1 负责生成 floor/wall 骨架，这是整个地图中**空间结构约束最强**的层次：

- 墙壁需要形成封闭或半封闭的房间边界；
- 走廊需要是连通的、宽度一致的通道；
- 整体拓扑（房间数、对称性、外围走廊）需要全局一致。

上述约束全部是**二维局部连通性约束**：一个墙壁格子是否合理，取决于它的上下左右四个邻格，而非它前后若干 token。一维位置编码使模型必须从数据中隐式学习这种行列边界，代价高昂且泛化差。

相比之下，Stage2/3 的任务（在走廊或房间内散布门/怪物/资源）对全局结构的空间一致性要求较低，位置编码的精确性影响较小。

---

## 改进方案

### 方案 A：二维因式分解位置嵌入（推荐首选）

#### 思路

将当前单一的一维位置嵌入替换为行嵌入与列嵌入的加和：

```
pos_embed[i, j] = row_embed[i] + col_embed[j]
```

这样，同一行的所有格子共享相同的行嵌入，同一列的所有格子共享相同的列嵌入。模型可以直接从嵌入中感知行列身份，而无需从一维序号中隐式推断。

#### 具体实现

```python
# 替换原有的 pos_embedding
self.row_embedding = nn.Parameter(torch.randn(1, MAP_H, d_model) * 0.02)
self.col_embedding = nn.Parameter(torch.randn(1, MAP_W, d_model) * 0.02)
```

前向传播中：

```python
# map: [B, H*W]
row_idx = torch.arange(MAP_H, device=map.device).repeat_interleave(MAP_W)
col_idx = torch.arange(MAP_W, device=map.device).repeat(MAP_H)
pos = self.row_embedding[0, row_idx] + self.col_embedding[0, col_idx]
x = self.tile_embedding(map) + pos.unsqueeze(0)
```

也可以预计算展开后的索引并缓存。

#### 特点

- 参数量变化：从 `169 × d_model` 变为 `13 × d_model + 13 × d_model = 26 × d_model`，显著减少，且参数共享有助于泛化；
- 无需修改注意力机制，改动最小；
- 直接赋予模型行列语义，改进效果立竿见影。

---

### 方案 B：二维相对位置偏置（推荐次选，与 A 叠加）

#### 思路

在注意力计算中，对每对 query-key 的打分加入一个可学习偏置，偏置由两个 token 的**相对行列偏移量**决定：

```
score(i, j) = (q_i · k_j) / sqrt(d) + B[Δrow, Δcol]
```

其中 `Δrow = row(i) - row(j)`，`Δcol = col(i) - col(j)`。偏置表 B 的形状为 `(2H-1, 2W-1)` = `(25, 25)`，每个注意力头各一张。

这种方式的核心优势：注意力打分天然理解"相邻格子应更强相关"，模型无需从位置嵌入中隐式学习距离感。

#### 具体实现

**步骤一：预计算相对位置索引表**

对于 13×13 的地图，预计算每对 token (i, j) 的相对位置，将二维偏移量映射为一维索引，供后续从偏置表中 gather：

```python
def build_relative_position_index(H: int, W: int) -> torch.Tensor:
    coords_h = torch.arange(H)
    coords_w = torch.arange(W)
    coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij')) # [2, H, W]
    coords_flat = coords.flatten(1) # [2, H*W]
    rel = coords_flat[:, :, None] - coords_flat[:, None, :] # [2, H*W, H*W]
    rel[0] += H - 1
    rel[1] += W - 1
    rel_index = rel[0] * (2 * W - 1) + rel[1] # [H*W, H*W]
    return rel_index
```

**步骤二：在模型中注册偏置表**

需要自定义 `SelfAttentionWithRPB`，替换 `nn.TransformerEncoderLayer` 中的注意力：

```python
class SelfAttentionWithRPB(nn.Module):
    def __init__(self, d_model: int, nhead: int, H: int, W: int):
        super().__init__()
        self.nhead = nhead
        self.d_head = d_model // nhead
        self.scale = self.d_head ** -0.5
        self.qkv = nn.Linear(d_model, d_model * 3)
        self.out_proj = nn.Linear(d_model, d_model)
        # 偏置表：(2H-1) * (2W-1) 个可能的相对位置，每个头各一组
        self.rel_bias_table = nn.Parameter(
            torch.zeros(nhead, (2 * H - 1) * (2 * W - 1))
        )
        rel_index = build_relative_position_index(H, W) # [H*W, H*W]
        self.register_buffer('rel_index', rel_index.flatten()) # [H*W * H*W]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.nhead, self.d_head)
        q, k, v = qkv.unbind(2) # 各 [B, N, nhead, d_head]
        q = q.permute(0, 2, 1, 3) # [B, nhead, N, d_head]
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)

        attn = (q @ k.transpose(-2, -1)) * self.scale # [B, nhead, N, N]

        # 从偏置表中取出对应偏置，reshape 成 [nhead, N, N]
        bias = self.rel_bias_table[:, self.rel_index].reshape(self.nhead, N, N)
        attn = attn + bias.unsqueeze(0)

        attn = attn.softmax(dim=-1)
        out = (attn @ v).permute(0, 2, 1, 3).reshape(B, N, C)
        return self.out_proj(out)
```

**步骤三：替换 TransformerEncoderLayer 中的注意力**

由于 PyTorch 标准 `TransformerEncoderLayer` 不支持直接替换注意力实现，需要手写包含 RPB 的编码器层：

```python
class RPBEncoderLayer(nn.Module):
    def __init__(self, d_model: int, nhead: int, dim_ff: int, H: int, W: int):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = SelfAttentionWithRPB(d_model, nhead, H, W)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, dim_ff),
            nn.GELU(),
            nn.Linear(dim_ff, d_model)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x
```

#### 参数量评估

对于 Stage1 MaskGIT（d_model=192, nhead=8, num_layers=6，H=W=13）：

- 偏置表：每层 `8 × 25 × 25 = 5000` 参数
- 共 6 层：30,000 参数
- 占总参数量的比例极低，但对 attention 的几何感知能力提升显著。

#### 特点

- 与方案 A 互补，可叠加使用；
- 理论上最接近"正确"的二维注意力感应偏置；
- 实现较复杂，需要手写注意力层；
- 参数量增加极少。

---

### 方案 C：轴向注意力（Axial Attention）

#### 思路

将每一个标准自注意力层替换为两个顺序执行的注意力：

1. **行轴注意力**：每一行内的格子互相注意，13 个 group，每组 13 个 token；
2. **列轴注意力**：每一列内的格子互相注意，13 个 group，每组 13 个 token。

两种注意力交替叠加：`Row → Col → Row → Col → ...`

```
标准自注意力：O((H*W)²) = O(169²) = 28,561
轴向注意力：  O(H * W²) + O(W * H²) = O(H*W*(H+W)) = O(169*26) ≈ 4,394
```

复杂度大幅下降，且强制模型以行/列为单位建立空间关联。

#### 具体实现

```python
class AxialAttention(nn.Module):
    def __init__(self, d_model: int, nhead: int, H: int, W: int):
        super().__init__()
        self.H = H
        self.W = W
        self.row_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.col_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.norm_row = nn.LayerNorm(d_model)
        self.norm_col = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape
        H, W = self.H, self.W
        x = x.reshape(B, H, W, C)

        # 行注意力：每行的 W 个 token 互相注意
        x_row = x.reshape(B * H, W, C)
        x_row, _ = self.row_attn(x_row, x_row, x_row)
        x = x + x_row.reshape(B, H, W, C)
        x = self.norm_row(x)

        # 列注意力：每列的 H 个 token 互相注意
        x_col = x.permute(0, 2, 1, 3).reshape(B * W, H, C)
        x_col, _ = self.col_attn(x_col, x_col, x_col)
        x = x + x_col.reshape(B, W, H, C).permute(0, 2, 1, 3)
        x = self.norm_col(x)

        return x.reshape(B, N, C)


class AxialEncoderLayer(nn.Module):
    def __init__(self, d_model: int, nhead: int, dim_ff: int, H: int, W: int):
        super().__init__()
        self.axial_attn = AxialAttention(d_model, nhead, H, W)
        self.norm_ff = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, dim_ff),
            nn.GELU(),
            nn.Linear(dim_ff, d_model)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.axial_attn(x)
        x = x + self.ffn(self.norm_ff(x))
        return x
```

#### 特点

- 最强的二维结构归纳偏置：明确区分行内关系和列内关系；
- 改造最彻底，需要替换整个 Transformer 编码器；
- 计算量低于标准全局注意力；
- 不适合 cross-attention 部分（decoder 中 map token cross-attend z 的部分保持不变）。

---

### 方案 D：双轴并行输入（Dual-Axis Input）

#### 思路

同一张地图展平两次，分别按行优先和列优先展平，两份序列并行送入各自的嵌入+编码器，最终在 token 维度相加合并：

```
地图 [H, W]
  ├─ 行优先展平 → [H*W] → Embedding + 行位置编码 → 编码器 A → [H*W, d_model]
  └─ 列优先展平 → [W*H] → Embedding + 列位置编码 → 编码器 B → [W*H, d_model]
       ↓ 重排到行优先顺序
  相加合并 → [H*W, d_model] → Decoder + z → logits
```

两个编码器可以共享权重（只有位置编码不同），以减少参数量。

#### 特点

- 无需修改注意力机制，复用现有 Transformer；
- 计算量加倍（两次编码），但可通过共享权重缓解；
- 同时为模型提供横向和纵向的序列上下文；
- 参数复用程度较高（共享编码器权重时）。

---

## 方案对比

| 方案                | 改动范围     | 额外参数量                         | 二维感知能力 | 实现复杂度 |
| ------------------- | ------------ | ---------------------------------- | ------------ | ---------- |
| A：二维因式位置嵌入 | 位置嵌入层   | 减少（共享）                       | 中等         | 低         |
| B：二维相对位置偏置 | 注意力层     | 极少（~30k）                       | 强           | 中等       |
| C：轴向注意力       | 整个编码器   | 基本不变                           | 最强         | 高         |
| D：双轴并行输入     | 输入与编码器 | 按编码器大小翻倍（共享权重则不变） | 中等         | 中等       |

---

## 推荐实施策略

### 第一步：替换位置嵌入（方案 A）

优先实施方案 A，这是改动最小、风险最低的基础改进。仅需修改 `GinkaMaskGIT` 和 `GinkaVQVAE` 中的 `pos_embedding` 初始化与使用方式。

修改范围：

- `ginka/maskGIT/model.py`：替换 `pos_embedding`，调整 `forward` 中的位置嵌入加法；
- `ginka/vqvae/model.py`：同上。

### 第二步：叠加二维相对位置偏置（方案 B）

在方案 A 的基础上，为 `GinkaMaskGIT` 的编码器部分叠加 RPB，只需新增一个自定义 `RPBEncoderLayer`，替换 `maskGIT.py` 中 `Transformer.encoder` 的层类型。

VQ-VAE 编码器的 RPB 改造优先级较低（VQ-VAE 负责全图压缩，对局部连通性感知需求低于 MaskGIT）。

### 可选第三步：轴向注意力替换（方案 C）

若前两步改进后 Stage1 质量仍不满足要求，可进一步将 `GinkaMaskGIT` 的编码器改为轴向注意力。由于 decoder（cross-attention 部分）不涉及空间 token，无需改动。

### 其他注意事项

- 上述改动仅针对 **Stage1 MaskGIT**。Stage2/3 可以同步修改，也可以保持原结构，视实际效果决定；
- 改动后需重置训练，位置嵌入的结构变化会使旧检查点不兼容（shape 不匹配）；
- 验证时重点观察生成地图中墙壁的连通性、房间闭合度和整体拓扑是否改善。