Dx's Tech Notes

框架图修改2.0

9 min read

框架图修改2.0

关联笔记：

目标

这份笔记整理的是两轮关于 AFB 流程图 _2 (cropped) (pdfresizer.com).pdf 的讨论结论。目标不是小修标注，而是按当前实际代码，把这张图从“概念示意图”改成“代码一致图”。

对应代码文件：

  • /Users/xinduan/FaF/ForensicsAdapter-main/model/adapters/BRA_trian_topk_CFM_adapter.py

核心对照位置：

  • FABiLevelRoutingAttention.forward()
  • BRAttnForAdapter.forward()

总体结论

当前这张图最大的问题不是美观，而是图的层级和代码层级没有对齐。现在的图把：

  • vision token 的 2D 重排
  • BRA 内部的 bi-level routing
  • token gathering
  • sparse attention
  • routing weights
  • 甚至 AFB 的整体效果可视化

混在了一张图里，但没有严格反映代码的先后顺序。

如果要对照代码大修，优先建议把这张图重定义为：

BRA inside AFB

也就是说，这张图主要讲：

  • 输入是 vision tokens
  • 经过 2D feature map 重排
  • 进行 region partition
  • 进行 qkv projection
  • 基于 Q_win/K_win 做 top-k routing
  • gather K/V
  • 做 sparse attention
  • 输出 updated visual tokens

不要再让它看起来像：

  • 原始 RGB 图像直接进入模块
  • 或者最终输出是“关注后的人脸图”

与代码不一致的关键问题

1. 输入画成了类似“原图”，但代码输入其实是 vision tokens

代码中真正进入 BRA 的是：

q_tok = x[:, :QL, :]
v_tok = x[:, QL:, :]
v_upd = self.bra_vis(v_tok, hw=(H, W))

也就是说：

  • q_tok 不进入 BRA 内部 routing
  • v_tok 才是 AFB/BRA 的输入

因此图左边不应再画成“图像块输入模块”，而应画成：

  • Visual Tokens
  • 或 Vision Tokens
  • shape 可标为 (B, HW, C)

2. 图里把 Avg Pool 放在最前面，这是不对的

代码顺序是：

qkv = self.qkv(x_win)
q, k, v = qkv.split([C, C, C], dim=-1)
q_win = q.mean(dim=(2, 3))
k_win = k.mean(dim=(2, 3))

正确顺序应是：

  1. 先 regionize
  2. 先 qkv projection
  3. 再对 Q^r/K^r 做 region mean

因此图里不应该是：

Input -> Avg Pool -> Q_win/K_win

而应该是：

Regionized Features -> Linear qkv Projection -> Q^r/K^r/V^r -> Region Mean on Q^r/K^r -> Q_win/K_win

3. 图中 Q 的含义不清，容易和 adapter 的 query token 混淆

当前图右侧有一个 Q 输入 sparse attention，但在全文里 query token 另有含义。

在代码里，BRA 内部的 Q 是：

  • 来自视觉特征的局部 query
  • 不是 adapter 的 learnable query token

因此图中建议改成：

  • Q^r
  • 或 Local Visual Queries
  • 或 Region Queries

4. 图里没有明确区分 I_win 和 W_win 的不同作用

代码中 routing 同时输出：

r_weight, r_idx = self.router(q_win, k_win)

即：

  • r_idx 对应 routing index
  • r_weight 对应 routing weights

后续：

  • r_idx 用来 gather K/V
  • r_weight 用来给 V_g 加权

所以图中必须清楚拆成：

  • A_win
  • Top-k Selection
  • I_win
  • W_win

而不是只画一个模糊的黄色 Routing Weights 框。

5. 图里输出不应该画成“人脸图”，而应该是 updated visual tokens

代码里 BRA 的输出是：

out = self.proj(out_2d).view(B, N, C)

也就是：

  • 输出仍然是 token 序列
  • 更准确地说是 updated visual tokens

所以图右侧建议改成：

  • Updated Visual Tokens
  • 或 v_upd

不要直接画成“关注后的人脸图”，那会让读者误以为这就是模块输出的显式可视化结果。

推荐的重画思路

方案选择

推荐把这张图定义为：

Detailed workflow of BRA inside the Artifact Focusing Block

这样图的范围清晰，避免把后续 query-to-vision attention 也硬塞进去。

如果以后要画完整 AFB block，可以另起一张图，在右侧补上：

  • q_tok
  • q_proj / k_proj / v_proj
  • query-to-vision attention
  • concat + out_proj

推荐的重画流程

下面是一版最贴代码的从左到右流程。

Step 0. 输入

左侧画一排或一列 token 小块：

  • Visual Tokens v_tok
  • shape: (B, HW, C)

这里不要用真实人脸图作为输入主体。

Step 1. Reshape to 2D Feature Map

箭头旁标：

  • reshape
  • 或 reshape to H × W

右边画一个规则网格，不画成人脸像素图，而画成抽象的 token map：

  • 2D Visual Feature Map
  • shape: (B, H, W, C)

说明：

  • 这是 token 的空间重排
  • 不是恢复原始 RGB 图像

Step 2. Partition into S × S Regions

在 2D feature map 上加粗分区线，切成 S × S 个大区块。

标注：

  • Partition into S × S Regions
  • S = 4
  • T = HW / S^2

旁边可以单独再画一个更简洁的区域化视图：

  • Regionized Features X^r

Step 3. Linear qkv Projection

在 regionized features 后放一个矩形框：

  • Linear qkv Projection

从该框分出三条支路：

  • Q^r
  • K^r
  • V^r

这一步对照代码：

qkv = self.qkv(x_win)
q, k, v = qkv.split([C, C, C], dim=-1)

Step 4. Region Mean on Q and K

在 Q^r 和 K^r 两条支路后各放一个小框：

  • Region Mean

输出为：

  • Q_win
  • K_win

V^r 不走这个平均分支，而是直接保留到后面的 gather/sparse attention。

Step 5. Region Routing

把 Q_win 和 K_win 接到一个 routing 模块中。

建议模块内部标明：

$$ A_{win} = \frac{Q_{win}K_{win}^{T}}{\sqrt{C}} $$

再从这个模块明确分成两步：

  • Top-k Selection
  • 输出 I_win
  • 输出 W_win

图里最好不要只写 Index，而是写得更完整：

  • Routing Indices I_win
  • Routing Weights W_win

Step 6. Token Gathering

单独一个框：

  • Token Gathering

输入：

  • K^r
  • V^r
  • I_win

输出：

  • K_g
  • V_g

这里一定要体现：

  • gather 的是 K/V
  • 不是 Q

Step 7. Weighting on Gathered Values

代码里有很重要的一步：

v_sel = v_sel * r_weight_expanded

这在图中必须体现出来。建议画成：

  • W_win
  • 乘到 V_g
  • 输出 Weighted V_g

否则图和代码会脱节。

Step 8. Sparse Attention

这个模块建议写成：

  • Sparse Attention

输入：

  • Q^r
  • K_g
  • Weighted V_g

不要只写一个模糊的 Q。

Step 9. Merge + LePE + Linear Projection

这部分建议不要拆得太碎，但至少要体现：

  • merge back to 2D
  • LePE
  • linear projection

可以画成一个组合模块：

  • Merge
  • DWConv-based LePE
  • Linear Projection

或者三个连着的小框。

Step 10. 输出

最终输出不要画成人脸图，而画成一排更新后的 token：

  • Updated Visual Tokens v_upd
  • shape: (B, HW, C)

关于 “Reshape to 2D Feature Map” 应该怎么画

最容易画错的地方，是把它画得像“token 变回了原始人脸图像”。

正确理解应该是：

  • 左边是 token sequence
  • 右边是带空间拓扑的 token grid

因此建议：

  • 左侧画 token blocks
  • 右侧画抽象网格图
  • 中间箭头标 reshape

不要把这一部分画成高还原度的人脸图片。

关于 “Linear qkv Projection” 应该怎么画

推荐画法：

画法 A：一个投影框，后面分三路

最推荐，最贴代码。

一个框：

  • Linear qkv Projection

后面三条支路：

  • Q^r
  • K^r
  • V^r

画法 B：三个并列线性框

也可以，但更占空间：

  • Linear W^q
  • Linear W^k
  • Linear W^v

对你这张图而言，建议用画法 A。

最终建议的框图文字版草图

可以直接照着下面这个顺序画：

  1. Visual Tokens v_tok
  2. Reshape to 2D Feature Map
  3. Partition into S × S Regions
  4. Linear qkv Projection
  5. Q^r / K^r / V^r
  6. Region Mean on Q^r and K^r
  7. Q_win / K_win
  8. Region Routing
  9. A_win -> Top-k -> I_win, W_win
  10. Token Gathering
  11. K_g / V_g
  12. W_win × V_g
  13. Sparse Attention
  14. Merge + LePE + Linear Projection
  15. Updated Visual Tokens v_upd

一句话总结

这张图最需要的不是补几个变量，而是把它从“概念图”改成“与代码严格一致的模块图”。

最关键的三点是：

  1. 输入是 vision tokens，不是原始图像
  2. 顺序是 reshape -> partition -> qkv -> region mean -> top-k routing -> gather -> sparse attention
  3. 输出是 updated visual tokens，不是可视化的人脸响应图

Graph View