如何在Flowformer中显式地添加attn_mask?

[Prot-debug]

作者您好，我对你们在 ICML 2022 的工作十分感兴趣，但在使用过程中有点疑问想请教一下，希望能够给我一些建议，谢谢。

1. 在代码中的 Flow_Attention_Causal 函数，casual 是通过 cumsum 函数实现的，但是在以往的注意力方法中都是通过 attn_mask 来进行显式地 mask，但使用乘法结合律的话需要先计算 KV，导致 attn_mask 的维度和 KV 维度不匹配，无法进行正常的 mask ，请问有没有什么方法依然可以显式地使用 attn_mask 进行 casual 操作？

谢谢你们精彩的工作，期待您的回答。

[wuhaixu2016]

您好，感谢关注！ 一般来说mask需要应对的情况有两种： （1）针对Causal任务进行上三角矩阵的mask：这个情况使用cumsum即可实现； （2）针对特定token进行mask：这种情况可以直接在v上进行mask，不必非要从attention上进行mask； 所以我觉得现在的设计已经可以满足主要情况的需求了，如果还有其他的特殊情况的话，可能需要额外的特殊设计，您也可以描述一下具体需求，再进行深入讨论。

[Prot-debug]

@wuhaixu2016 您好，

针对您说的第二种情况，如果我在代码中想使用 cross attention，假设我的 query 维度是 m×d（忽略batch维度），key 和 value 的维度是 n×d 。并且，我想在做 cross attention 的时候，mask 掉 value 中的部分token。

1. 如果直接使用 Flow_Attention_Causal 函数的话，因为 m≠n，会导致维度不匹配报错，所以我使用了 Flow_Attention 函数，但在使用时我有一个疑问：您在上面的回答中说，可以直接在 v 上进行 mask，这种情况我是用 0 来 mask，还是用 -inf 来 mask 呢？这个 mask 是不是应该连同 k 和 v 一起 mask 掉，还是只用 mask v 就可以了？
2. 在以往的 attention weight 中，在 $$q^T*k$$ 以后，我可以给部分 token 的位置加上 log(x) 来显式地修改这部分 token 在 softmax 中的占比，但在 flowformer 中，还有办法实现这样的操作吗？
3. 另外，如果完成了上述的操作的前提下，使用 Flow_Attention 函数进行 cross attention 操作会导致信息泄漏吗？

不知道上面的描述是否清晰，希望您有时间的时候可以回复一下，谢谢！

[wuhaixu2016]

1. 如果你想达到一个上三角的mask，实现causal，只能使用casual版本的flowformer。如果因为q和k长度不一致，可以短的一方用0补齐
2. 用0来mask
3. softmax会造成二次复杂度，所以不能有。您可以在flowformer的softmax激活函数那里加上temperature来调整attention的peaky程度
4. 请看1，如果直接mask v会造成泄漏。因此，如果您想实现causal，那就只能把q和k长度补齐，然后使用causal 版本的flow attention

[Prot-debug]

如果我在做 cross attention 的时候，想实现对 memory（value和key） 的 mask，且这个 mask 矩阵并不是上三角的，而是对部分特定的 token 进行 mask。比如说 query 的维度是 m×d，一共有 m 个 token，key 和 value 的维度是 n×d，一共有 n 个 token，我在做第一个query的时候，只想对第一个和第二个 memory token 进行mask，在第二个query的时候，只想对第二个和第三个 memory token 进行mask，以此类推，请问这种情况有办法做到吗？

[wuhaixu2016]

这个操作有点复杂，我目前也没有想到解法。 而且目前我能想到的其他的efficient transformer也无法完成这个要求，可能只能用标准attention，计算map之后来进行操作了。

[Prot-debug]

那如果考虑一种简单一点的情况，比如说，我在做 cross attention 的时候，不对 memory（value和key） 进行 mask，也不对 query 进行 self attention 操作，而是直接使用 Flow_Attention 函数用 query 对 memory 进行查询。这种情况下，会有信息泄漏的风险吗？如果有，是在什么地方泄漏的呢？

[wuhaixu2016]

按照我的理解，是会泄漏的，因为在固定源和求汇的时候，源和的计算会泄漏未来信息。

[Prot-debug]

那如果我这里的 memory 是全部可访问的 token（也就是不包含未来信息），只有 query 里才包含未来信息，这种情况是不是也会发生泄漏呢？

[wuhaixu2016]

如果直接计算会存在泄漏。 不过这个应该是完全可以实现的，对query的计算使用cusum，对于key的计算使用正常的sum即可。

[Prot-debug]

好的，谢谢你耐心且高效的解答。后期如果我的论文能顺利发表的话，我会引用你们的工作，谢谢！

[wuhaixu2016]

感谢关注！