Sliding Tile Attentionについて

以下の論文を読む。僅かな性能低下でHunyuanVideoのAttn高速化が述べられている。
論文にあるSliding Tile AttentionのAttention mapを描いてみたい。

https://www.arxiv.org/pdf/2502.04507


# 3D Attentionの冗長性
Videoを作るHunyuanVideoモデルでは動画時間(生成フレーム)が長くなるとAttentionの計算時間がtoken数に対して2次で増えるのでAttentionの計算時間が支配的になる。例えば10kのtokenの計算時間を$100$とした時、115kの計算時間は$5400.4$となり、この5400.4のうちFFNが$43.2\cdot \frac{115}{10}=496.8$、QKVが$21.6\cdot \frac{115}{10}=248.4$、ATTNが$35.2\cdot (\frac{115}{10})^2=4655.2$となる。
ここで10kとはlatent次元で50x50x4(動画次元で400x400x17)、115kとはlatent次元で60x80x24(動画次元で480x640x97)であろう。

しかし、実のところ時間的や空間的に隣接するToken間でのAttention項は意味があるが、遠く離れたToken間でのAttentionを計算してもほとんど意味をなさない。
要するにText2Videoというタスクにおいてはこれらの遠く離れたAttention計算は計算してもその影響度はほとんどゼロに近くこれを省略しようが、性能の劣化はあまりないらしい。




![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/ff50965a-dd24-4b57-91ce-f4f258f501ac.png)

token_numに対してFFNとQKVは一次関数的に増え、ATTNは二次関数的に増えるとし、10kの時の計算量を100とする。動画のサイズがある程度以上になるとATTN計算が支配的になる。

|token_num |FFN  |QKV  |ATTN  |合計  |latent_size  |videosize |
|---|---|---|---|---|---|---|
|10k|43.2<br>(43.2%)  |21.6<br>(21.6%)  |35.2<br>(35.2%)  |100  |50x50x4  |400x400x17  |
|15k|64.8<br>(36.7%)  |32.4<br>(18.4%)  |79.2<br>(44.9%)  |176.4  |50x75x4  |400x600x17  |
|19.2k|82.9<br>(32.6%)  |41.5<br>(16.3%)  |129.8<br>(51.1%)  |254.2  |60x80x4  |480x640x17 |
|20k|86.4<br>(32.0%)  |43.2<br>(16.0%)  |140.8<br>(52.1%)  |270.4  |50x50x8  |400x400x33 |
|38.4k|165.9<br>(21.6%)  |82.9<br>(10.8%)  |519.0<br>(67.6%)  |767.9  |60x80x8  |480x640x33 |
|57.6k|248.8<br>(16.1%)  |124.4<br>(8.1%)  |1167.9<br>(75.8%)  |1541.1  |60x80x12  |480x640x49 |
|76.8k|331.8<br>(12.9%)  |165.9<br>(6.4%)  |2076.2<br>(80.7%)  |2573.8  |60x80x16  |480x640x65 |
|115k|496.8<br>(9.2%)  |248.4<br>(4.6%)  |4655.2<br>(86.2%)  |5400.4  |60x80x24  |480x640x97  |
|153.6k|663.6<br>(7.1%)  |331.8<br>(3.6%)  |8304.7<br>(89.3%)  |9300.0  |60x80x32  |480x640x129  |

https://arxiv.org/pdf/2204.07143

https://arxiv.org/pdf/2406.08552

https://arxiv.org/pdf/2412.16112

# NATTEN
このようなAttentionの冗長性を取り除くのにNeighborhood Attention(NATTEN)があり、以下のような図が知られている。しかし、この図を見ただけでは背景がよく分からないのでどうやってこの図を描いたのかがいま一つピンとこない。プログラム的にこの図をどうやって得ているのかを考えたい。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/08e974bd-bdc7-4f3b-8c10-e0fe4a83765e.png)

まず最初に画像サイズ(latentサイズ)が24x24とし、Attentionが有効となる距離を12x12とする。以下のようなプログラムを書けばNATTENのマップは書ける。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/9198d4e0-9bd4-4c52-8745-21010be6be73.png)



```python
import numpy as np
import matplotlib.pyplot as plt

index = np.arange(24*24).reshape(24,24)
map = np.zeros((24*24,24*24))

for i in range(24*24):
    x, y = i%24, i//24
    x = np.clip(x, 6, 18)
    y = np.clip(y, 6, 18)
    window_index = index[y-6:y+6,x-6:x+6].flatten()
    map[i,window_index] = 1.0

u = np.arange(24*24)
v = np.arange(24*24)
X, Y = np.meshgrid(u, v)
plt.pcolormesh(X,Y,map)
plt.show()
```

# Morton curve

次にFlexAttentionの図を見るとMorton Curveなる単語が見られる。
これは上述プログラムのindexカウントを改良するためのものだろう。現行のindexは画像端まで数えて折り返してるのでblock_window上のAttention範囲を考えると数字の範囲が飛びがちである。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/b5b12cef-5956-4e0b-a4bf-93a9120ba6b9.png)

参考までにspace-filling curveの資料を示しておくがZ-order_curve(Morton)は24x24を``index = np.arange(24*24).reshape(3,3,2,2,2,2,2,2).transpose(0,2,4,6,1,3,5,7).reshape(24,24)``とするのに等しい。このように二次元空間的indexと空間距離を縮める並べ方の工夫について調べたが、**別にTokenの並べ方で解決している訳ではなさそうである。**(参考程度)

```python
import numpy as np
index = np.arange(24*24).reshape(3,3,2,2,2,2,2,2).transpose(0,2,4,6,1,3,5,7).reshape(24,24)
print(index)
--------------------------------------------------
[[  0   1   4   5  16  17  20  21  64  65  68  69  80  81  84  85 128 129
  132 133 144 145 148 149]
 [  2   3   6   7  18  19  22  23  66  67  70  71  82  83  86  87 130 131
  134 135 146 147 150 151]
 [  8   9  12  13  24  25  28  29  72  73  76  77  88  89  92  93 136 137
  140 141 152 153 156 157]
 [ 10  11  14  15  26  27  30  31  74  75  78  79  90  91  94  95 138 139
  142 143 154 155 158 159]
 [ 32  33  36  37  48  49  52  53  96  97 100 101 112 113 116 117 160 161
  164 165 176 177 180 181]
 [ 34  35  38  39  50  51  54  55  98  99 102 103 114 115 118 119 162 163
  166 167 178 179 182 183]
 [ 40  41  44  45  56  57  60  61 104 105 108 109 120 121 124 125 168 169
  172 173 184 185 188 189]
 [ 42  43  46  47  58  59  62  63 106 107 110 111 122 123 126 127 170 171
  174 175 186 187 190 191]
  ...
```


https://en.wikipedia.org/wiki/Z-order_curve

https://en.wikipedia.org/wiki/Peano_curve

https://en.wikipedia.org/wiki/Sierpi%C5%84ski_curve

https://en.wikipedia.org/wiki/Moore_curve

https://en.wikipedia.org/wiki/Hilbert_curve


このZ-order_curve(Morton curve)の時のAttention Maskを参考までに描けば、以下のようになる。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/0c1e8c85-dc36-4db7-970d-0f25750ba348.png)


```python
import numpy as np
import matplotlib.pyplot as plt

index = np.arange(24*24).reshape(3,3,2,2,2,2,2,2).transpose(0,2,4,6,1,3,5,7).reshape(24,24)
map = np.zeros((24*24,24*24))

for i in range(24*24):
    min_index = np.argmin(np.abs(index-i))
    x = min_index%24
    y = min_index//24
    x = np.clip(x, 6, 18)
    y = np.clip(y, 6, 18)
    window_index = index[y-6:y+6,x-6:x+6].flatten()
    map[i,window_index] = 1.0

u = np.arange(24*24)
v = np.arange(24*24)
X, Y = np.meshgrid(u, v)
plt.pcolormesh(X,Y,map)
plt.show()
```


https://arxiv.org/pdf/2412.05496

https://openreview.net/pdf?id=MdWqXRaKdp


# Tiled NATTEN
とりあえずindexを4x4のタイルを考え、その4x4タイルを6x6で敷き詰めた図で12x12のblock_windowを考えたmapが以下である。この図がTiled NATTENと等しいのが確認できる。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/a08a28b0-27a4-4d78-9917-c5b28f9a42f6.png)

```python
import numpy as np
import matplotlib.pyplot as plt

index = np.arange(24*24).reshape(6,6,4,4).transpose(0,2,1,3).reshape(24,24)
map = np.zeros((24*24,24*24))

for i in range(24*24):
    block_index = i//16
    block_inner_index = i%16
    inner_x, inner_y = block_inner_index%4, block_inner_index//4
    block_x, block_y = 4*(block_index%6), 4*(block_index//6)
    x = block_x + inner_x
    y = block_y + inner_y
    x = np.clip(x, 6, 18)
    y = np.clip(y, 6, 18)
    window_index = index[y-6:y+6,x-6:x+6].flatten()
    map[i,window_index] = 1.0

u = np.arange(24*24)
v = np.arange(24*24)
X, Y = np.meshgrid(u, v)
plt.pcolormesh(X,Y,map)
plt.show()
```



# STA(Sliding Tile Attention)
ここまで考えたらSTAが何なのかは何となく理解してきた。要するに前述のコードのinner_xとinner_yを無視すれば綺麗にAttention領域をblock状にして計算できる筈だ。(4以下のwindow幅への正確性はある程度減るが)。適当に4の倍数window取り出し範囲を-4～+8に変えれば以下のような図が描ける。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/ae0e50c3-bd3a-4c06-8cef-effdd99ba046.png)


```python
import numpy as np
import matplotlib.pyplot as plt

index = np.arange(24*24).reshape(6,6,4,4).transpose(0,2,1,3).reshape(24,24)
map = np.zeros((24*24,24*24))

for i in range(24*24):
    block_index = i//16
    block_x, block_y = 4*(block_index%6), 4*(block_index//6)
    x = block_x
    y = block_y
    x = np.clip(x, 4, 16)
    y = np.clip(y, 4, 16)
    window_index = index[y-4:y+8,x-4:x+8].flatten()
    map[i,window_index] = 1.0

u = np.arange(24*24)
v = np.arange(24*24)
X, Y = np.meshgrid(u, v)
plt.pcolormesh(X,Y,map)
plt.show()
```

# 多段STA
4x4のタイルが3x3に並んでいてこれが更に2x2で並んでいるとする。
周期的ではないが多少無視すれば簡単な図形に近似しやすく見える。後述のpatchfyはこれに近いのかと思う。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/cb9cc9c1-af02-4b70-94b3-af2753cdb80e.png)

```python
import numpy as np
import matplotlib.pyplot as plt

index = np.arange(24*24).reshape(2,2,3,3,4,4).transpose(0,2,4,1,3,5).reshape(24,24)
map = np.zeros((24*24,24*24))

for i in range(24*24):
    min_index = np.argmin(np.abs(index-i))
    x = min_index%24
    y = min_index//24
    x = 4*(x//4)
    y = 4*(y//4)
    x = np.clip(x, 4, 16)
    y = np.clip(y, 4, 16)
    window_index = index[y-4:y+8,x-4:x+8].flatten()
    map[i,window_index] = 1.0

u = np.arange(24*24)
v = np.arange(24*24)
X, Y = np.meshgrid(u, v)
plt.pcolormesh(X,Y,map)
plt.show()
```

# Hilbert curve
Hilbert curveでAttention Mapを考えてみる。Hilbeltは$2^N$である必要があり、24x24を考えることが出来ないのでここでは32x32のindexの並び順を操作する。windowサイズを16とした時、このコードは以下のように書ける。STAのようにここでは4以下の空間座標を省略する。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/d5ddffe8-d47b-4460-ae7f-412bdd3513ea.png)



```python
import numpy as np
import matplotlib.pyplot as plt
from hilbertcurve.hilbertcurve import HilbertCurve

index = np.arange(32*32).reshape(32,32)
hilbert_curve = HilbertCurve(5, 2)
distances = list(range(32*32))
points = hilbert_curve.points_from_distances(distances)
for point, dist in zip(points, distances):
    index[point[0], point[1]] = int(dist)
print(index)
map = np.zeros((32*32,32*32))

for i in range(32*32):
    min_index = np.argmin(np.abs(index-i))
    x = min_index%32
    y = min_index//32
    x = 4*(x//4)
    y = 4*(y//4)
    x = np.clip(x, 8, 24)
    y = np.clip(y, 8, 24)
    window_index = index[y-8:y+8,x-8:x+8].flatten()
    map[i,window_index] = 1.0

u = np.arange(32*32)
v = np.arange(32*32)
X, Y = np.meshgrid(u, v)
plt.pcolormesh(X,Y,map)
plt.show()
-----------------------------------
[[   0    1   14 ...  339  340  341]
 [   3    2   13 ...  338  343  342]
 [   4    7    8 ...  349  344  345]
 ...
 [1019 1016 1015 ...  674  679  678]
 [1020 1021 1010 ...  685  680  681]
 [1023 1022 1009 ...  684  683  682]]
```

``x = 4*(x//4),y = 4*(y//4)``をコメントアウトすれば以下である。フラクタル的なAttentionMapが見れる。このHilbert curveは対角近傍のAttentionマップに近いので**入力Tokenの適切な並べ替えのみ**で対角近傍計算でよい効率的なAttentionMapを生成する道もあるかもしれない。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/970456de-1045-43f9-ae62-ef61453ae92c.png)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/92cc31a4-5e2a-49a3-a91a-b293f0f50b62.png)



https://github.com/galtay/hilbertcurve



# patchfy(Lumina-video)
複雑なAttentionマップを考えずともLumina-videoではもっと簡単なpatchfyでAttentionの及ぶ長さを調整しているのかと思う。多分だけれども。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/9a21cf8e-1238-45cd-8dc2-20cf093f1fc6.png)
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/f1fb61f7-e914-4a06-ad52-a677ff0271bd.png)



https://arxiv.org/pdf/2502.06782

# 推論stepによる依存性
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/85a638bb-6b2f-4abb-8d08-9b90993551fb.png)![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/fb055966-638b-4d01-a9db-e2e4468d4e57.png)


Self-Attention(spatial,temporal Attention)の重要度は推論step(t)においても変動するのでこれも参考に設定を変えるべきなのかもしれない。


https://openreview.net/pdf?id=hDBrQ4DApF

https://arxiv.org/pdf/2411.19108

# その他参考：

https://www.reddit.com/r/StableDiffusion/comments/1itdpd4/sliding_tile_attention_a_new_method_that_speeds/

https://hao-ai-lab.github.io/blogs/sta/

# まとめ
動画生成モデルの論文を眺めた時、その生成のAttn計算を高速化する論文を幾つか見た。
FlashAttnetion2,3とかsdpa(Scaled Dot-Product Attention)とかsage attentionのAttention自体の高速化とは別口で、Attentionの計算領域を空間的や時間的に制限してそれを抑えるようだ。

その論文にある図が自分には直感的には理解できなかったので理解するためにAttntion Mapを自分で描いてみた。