多模态模型如何处理任意分辨率输入——Tiling与Packing技术详解(part2)

接上文http://myhz0606.com/article/vllm_any_resolution，补充一些公式推导

1 为什么block-diagonal attention mask能避免packing时attention跨图片计算

给定多张图片在sequence维度上拼接的序列

$$x=\left[\begin{array}{c}{p}_1\\ p_2\\ ⋮\\ p_n\end{array}\right],p_i\in {\mathbb{R}}^{L_i\times d},L=\sum _{i=1}^n{L}_i,x\in {\mathbb{R}}^{L\times d}\text{\hspace{1em}(1)}$$

给定Block Diagnoal Attention Mask

$$\begin{array}{rlr}\mathcal{M}& =\left[\begin{array}{ccccc}& m_{11}& m_{12}& \cdots & m_{1n}\\ & m_{21}& m_{22}& \cdots & m_{2n}\\ & ⋮& ⋮& \ddots & ⋮\\ & m_{n1}& m_{n2}& \cdots & m_{nn}\end{array}\right]& \\ m_{ij}& =\{\begin{array}{ll}{0}_{L_i\times L_i}& \mathrm{if}i=j\\ -\inf \cdot 1_{L_i\times L_j}& \mathrm{if}i\ne j\end{array}& \end{array}\text{\hspace{1em}(2)}$$

下面来计算Attention

先通过linear计算QKV，线性层时channel-mix的运算，没有token间的信息交互

$$\begin{array}{rl}Q& =\mathrm{Linear}(x)=\left[\begin{array}{c}{q}_1\\ ⋮\\ q_n\end{array}\right],K=\mathrm{Linear}(x)=\left[\begin{array}{c}{k}_1\\ ⋮\\ k_n\end{array}\right],V=\mathrm{Linear}(x)=\left[\begin{array}{c}{v}_1\\ ⋮\\ v_n\end{array}\right]\end{array}\text{\hspace{1em}(3)}$$

再有Block Diagnoal Attention Mask 的Attention计算

$$\begin{array}{rlr}\mathrm{Attention}(x)& =\mathrm{Softmax}(\frac{Q{K}^T}{\sqrt{d}}+\mathcal{M})V& \\ & =\mathrm{Softmax}(\frac{1}{\sqrt{d}}\left[\begin{array}{c}{q}_1\\ ⋮\\ q_n\end{array}\right]\left[\begin{array}{ccc}{k}_1^T& \cdots & k_n^T\end{array}\right]+\left[\begin{array}{cccc}& m_{11}& \cdots & m_{1n}\\ & ⋮& \ddots & ⋮\\ & m_{n1}& \cdots & m_{nn}\end{array}\right])\left[\begin{array}{c}{v}_1\\ ⋮\\ v_n\end{array}\right]& \\ & =\mathrm{Softmax}\left(\frac{1}{\sqrt{d}}\left[\begin{array}{cccc}& q_1{k}_1^T+m_{11}& \cdots & q_1{k}_n^T+m_{1n}\\ & ⋮& \ddots & ⋮\\ & q_n{k}_1^T+m_{11}& \cdots & q_n{k}_n^T+m_{nn}\end{array}\right]\right)\left[\begin{array}{c}{v}_1\\ ⋮\\ v_n\end{array}\right]& \text{because. }\frac{\mathcal{M}}{\sqrt{d}}=\mathcal{M}\\ & =\left[\begin{array}{cccc}& \frac{\exp \{(q_1{k}_1^T+m_{11})/\sqrt{d}\}}{{\sum }_{j=1}^L\exp \{(q_1{k}_j^T+m_{1j})/\sqrt{d}\}}& \cdots & \frac{\exp \{(q_1{k}_n^T+m_{1n})/\sqrt{d}\}}{{\sum }_{j=1}^L\exp \{(q_1{k}_j^T+m_{1j})/\sqrt{d}\}}\\ & ⋮& \ddots & ⋮\\ & \frac{\exp \{(q_n{k}_1^T+m_{n1})/\sqrt{d}\}}{{\sum }_{j=1}^L\exp \{(q_n{k}_j^T+m_{nj})/\sqrt{d}\}}& \cdots & \frac{\exp \{(q_n{k}_n^T+m_{nn})/\sqrt{d}\}}{{\sum }_{j=1}^L\exp \{(q_n{k}_j^T+m_{nj})/\sqrt{d}\}}\end{array}\right]\left[\begin{array}{c}{v}_1\\ ⋮\\ v_n\end{array}\right]& \\ & =\left[\begin{array}{cccc}& \frac{\exp \{(q_1{k}_1^T)/\sqrt{d}\}}{{\sum }_{j=1}^{L_1}\exp \{(q_1{k}_j^T)/\sqrt{d}\}}& \cdots & 0\\ & ⋮& \ddots & ⋮\\ & 0& \cdots & \frac{\exp \{(q_n{k}_n^T)/\sqrt{d}\}}{{\sum }_{j=L-L_n}^L\exp \{(q_n{k}_j^T)/\sqrt{d}\}}\end{array}\right]\left[\begin{array}{c}{v}_1\\ ⋮\\ v_n\end{array}\right]& \exp 0=0,\exp -\inf =0\\ & =\left[\begin{array}{c}\frac{\exp \{(q_1{k}_1^T)/\sqrt{d}\}}{{\sum }_{j=1}^{L_1}\exp \{(q_1{k}_j^T)/\sqrt{d}\}}{v}_1\\ ⋮\\ \frac{\exp \{(q_n{k}_n^T)/\sqrt{d}\}}{{\sum }_{j=L-L_n}^L\exp \{(q_n{k}_j^T)/\sqrt{d}\}}{v}_n\end{array}\right]& \\ & =\left[\begin{array}{c}\mathrm{Attention}(p_1)\\ \mathrm{Attention}(p_2)\\ ⋮\\ \mathrm{Attention}(p_n)\end{array}\right]& \\ & & \end{array}\text{\hspace{1em}(4)}$$

从上述推导可见Block Diagnoal Attention Mask 能有效避免跨图片的token交互。

2 为什么Cross-Attention融合要引入attention mask

在cross-attention融合方案中（上文3.2节）。文本特征作为attention的query，图片特征作为key，value。由于在动态分辨率batch推理时，每张图片所得到的图片特征长度是不同的，因此在组batch时，需要进行padding补齐。为了避免padding token干扰attention score，需引入attention mask。

下面以batch中的一个sample为例，说明attention mask的工作机制：

Q为文本特征，shape为$M\times d$

K， V为图片特征, shape为 $L\times d$ ， 其中slice为$[1:L_1]$为非padding的特征

$$\begin{array}{rl}Q& =\left[\begin{array}{c}{q}_1\\ q_2\\ ⋮\\ q_M\end{array}\right]K=\left[\begin{array}{c}{k}_1\\ k_2\\ ⋮\\ k_L\end{array}\right],V=\left[\begin{array}{c}{v}_1\\ v_2\\ ⋮\\ v_L\end{array}\right]\end{array}\text{\hspace{1em}(5)}$$

为了保证padding位置特征不会干扰attention score，需引入的mask如下

$$\begin{array}{rl}\mathcal{M}& =\left[\begin{array}{c}\stackrel{\#L_1}{\overbrace{0,0,\cdots ,0}},\stackrel{\#L_2}{\overbrace{-\inf ,\cdots ,-\inf }}\\ 0,0,\cdots ,0,-\inf ,\cdots ,-\inf \\ \cdots \\ 0,0,\cdots ,0,-\inf ,\cdots ,-\inf \end{array}\right]\in {\mathbb{R}}^{M\times L}\\ & \\ m& =\left[\begin{array}{c}\stackrel{\#L_1}{\overbrace{0,0,\cdots ,0}},\stackrel{\#L_2}{\overbrace{-\inf ,\cdots ,-\inf }}\end{array}\right]\end{array}\text{\hspace{1em}(6)}$$

证明：

$$\begin{array}{rl}\mathrm{Attention}(Q,K,V)=\left[\begin{array}{c}{y}_1\\ y_2\\ ⋮\\ y_M\end{array}\right]& =\mathrm{Softmax}(\frac{Q{K}^T}{\sqrt{d}}+\mathcal{M})V\\ & =\mathrm{Softmax}(\frac{1}{\sqrt{d}}\left[\begin{array}{c}{q}_1\\ q_2\\ ⋮\\ q_M\end{array}\right]\left[\begin{array}{c}{k}_1^T,k_2^T,\cdots ,k_L^T\end{array}\right]+\mathcal{M})\left[\begin{array}{c}{v}_1\\ v_2\\ ⋮\\ v_L\end{array}\right]\end{array}\text{\hspace{1em}(7)}$$

以$y_t$为视角

yt=∑i=1LSoftmax(qtkiTd+m)vi=∑i=1Lexp⁡{qtkiTd+mi}∑i=1Lexp⁡{qtkiTd+mi}vi=1∑i=1L1exp⁡{qtkiTd+0}+∑i=L1Lexp⁡{qtkiTd−inf⁡}⏟=0[∑i=1L1exp⁡{qtkiTd+0}vi+∑i=L1Lexp⁡{qtkiTd−inf⁡}vi⏟=0]=1∑i=1L1exp⁡{qtkiTd+0}[∑i=1L1exp⁡{qtkiTd+0}vi]=∑i=1L1Softmax(qtkiTd)vi](8) \begin{aligned}y_t &= \sum_{i=1}^{L} \mathrm{Softmax} (\frac{q_t k_i^T}{\sqrt{d}} + m)v_i \\&= \sum_{i=1}^{L}\frac{\exp\{{\frac{q_t k_i^T }{\sqrt{d}}} + m_i \}}{\sum_{i=1}^{L} \exp \{\frac{q_t k_i^T}{\sqrt{d}} + m_i \}} v_i \\&= \frac{1}{{\sum_{i=1}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}}+0\}}+\underbrace{{\sum_{i=L_1}^{L} \exp \{\frac{q_t k_i^T}{\sqrt{d}}-\mathrm{\inf}\}}}_{=0}}\left[\sum_{i=1}^{L_1}\exp\{{\frac{q_t k_i^T }{\sqrt{d}}} + 0 \} v_i+ \underbrace{\sum_{i=L_1}^{L}\exp\{{\frac{q_t k_i^T }{\sqrt{d}}} -\mathrm{\inf} \} v_i}_{=0}\right] \\&= \frac{1}{{\sum_{i=1}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}}+0\}}}\left[\sum_{i=1}^{L_1}\exp\{{\frac{q_t k_i^T }{\sqrt{d}}} + 0 \} v_i\right] \\&= \sum_{i=1}^{L_1} \mathrm{Softmax} (\frac{q_t k_i^T}{\sqrt{d}})v_i \end{aligned} ] \tag{8} yt​​=i=1∑L​Softmax(d​qt​kiT​​+m)vi​=i=1∑L​∑i=1L​exp{d​qt​kiT​​+mi​}exp{d​qt​kiT​​+mi​}​vi​=∑i=1L1​​exp{d​qt​kiT​​+0}+=0i=L1​∑L​exp{d​qt​kiT​​−inf}​​1​​i=1∑L1​​exp{d​qt​kiT​​+0}vi​+=0i=L1​∑L​exp{d​qt​kiT​​−inf}vi​​​​=∑i=1L1​​exp{d​qt​kiT​​+0}1​[i=1∑L1​​exp{d​qt​kiT​​+0}vi​]=i=1∑L1​​Softmax(d​qt​kiT​​)vi​​](8)

可见，引入式6的mask后使得：$y_t={\sum }_{i=1}^L\mathrm{Softmax}(\frac{q_t{k}_i^T}{\sqrt{d}}+m)v_i={\sum }_{i=1}^{L_1}\mathrm{Softmax}(\frac{q_t{k}_i^T}{\sqrt{d}})v_i$，满足了计算等价。

再看看如果不加mask，结果有何差异：

不妨将不加mask的$y_t$记作$y_t^{\prime }$。以padding zero为例，即$K=[\underset{\#L_1}{\underbrace{k_1,\cdots ,k_{L_1}}},\underset{\#L_2}{\underbrace{0,\cdots ,0}}{]}^T,V=[\underset{\#L_1}{\underbrace{v_1,\cdots ,v_{L_1}}},\underset{\#L_2}{\underbrace{0,\cdots ,0}}{]}^T$

yt′=∑i=1LSoftmax(qtkiTd)vi=1∑i=1Lexp⁡{qtkiTd}∑i=1Lexp⁡{qtkiTd}vi=1∑i=1L1exp⁡{qtkiTd}+∑i=L1Lexp⁡{qtkiT⏞=0d⏟=L−L1=L2}[∑i=1L1exp⁡{qtkiTd}vi+∑i=L1Lexp⁡{qtkiTd}vi⏟=0]=1∑i=0L1exp⁡{qtkiTd}+L2∑i=1L1exp⁡{qtkiTd}vi=∑i=1L1exp⁡{qtkiTd}∑i=1L1exp⁡{qtkiTd}+L2∑i=0L1Softmax(qtkiTd)vi=∑i=1L1exp⁡{qtkiTd}∑i=1L1exp⁡{qtkiTd}+L2yt(9) \begin{aligned}y_t' &= \sum_{i=1}^{L} \mathrm{Softmax} (\frac{q_t k_i^T}{\sqrt{d}})v_i \\&= \frac{1}{{\sum_{i=1}^{L} \exp \{\frac{q_t k_i^T}{\sqrt{d}} \}} }\sum_{i=1}^{L}\exp\{{\frac{q_t k_i^T }{\sqrt{d}}} \}v_i\\&= \frac{1}{ \sum_{i=1}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}} \} + \underbrace{ \sum_{i=L_1}^{L} \exp \{\frac{\overbrace{ q_t k_i^T}^{=0}}{\sqrt{d}} }_{=L-L_{1}=L_2} \} }\left [\sum_{i=1}^{L_1}\exp\{{\frac{q_t k_i^T }{\sqrt{d}}} \}v_i + \underbrace{ \sum_{i=L_1}^{L}\exp\{{\frac{q_t k_i^T }{\sqrt{d}}} \}v_i }_{=0}\right] \\&= \frac{1}{\sum_{i=0}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}} \} + L_{2}}\sum_{i=1}^{L_1}\exp\{{\frac{q_t k_i^T }{\sqrt{d}}} \}v_i \\&= \frac{\sum_{i=1}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}} \} }{\sum_{i=1}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}} \} + L_{2}} \sum_{i=0}^{L_1} \mathrm{Softmax} (\frac{q_t k_i^T}{\sqrt{d}})v_i \\ &= \frac{\sum_{i=1}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}} \} }{\sum_{i=1}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}} \} + L_{2}}y_t\end{aligned} \tag{9} yt′​​=i=1∑L​Softmax(d​qt​kiT​​)vi​=∑i=1L​exp{d​qt​kiT​​}1​i=1∑L​exp{d​qt​kiT​​}vi​=∑i=1L1​​exp{d​qt​kiT​​}+=L−L1​=L2​i=L1​∑L​exp{d​qt​kiT​​=0​​​​}1​​i=1∑L1​​exp{d​qt​kiT​​}vi​+=0i=L1​∑L​exp{d​qt​kiT​​}vi​​​​=∑i=0L1​​exp{d​qt​kiT​​}+L2​1​i=1∑L1​​exp{d​qt​kiT​​}vi​=∑i=1L1​​exp{d​qt​kiT​​}+L2​∑i=1L1​​exp{d​qt​kiT​​}​i=0∑L1​​Softmax(d​qt​kiT​​)vi​=∑i=1L1​​exp{d​qt​kiT​​}+L2​∑i=1L1​​exp{d​qt​kiT​​}​yt​​(9)

从推导结果可见，二者存在一个∑i=1L1exp⁡{qtkiTd}∑i=1L1exp⁡{qtkiTd}+L2\frac{\sum_{i=1}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}} \} }{\sum_{i=1}^{L_1} \exp \{\frac{q_t k_i^T}{\sqrt{d}} \} + L_{2}}∑i=1L1​​exp{d​qt​kiT​​}+L2​∑i=1L1​​exp{d​qt​kiT​​}​系数差异。