Mean PoolingとAttention Poolingなどについて

背景

これまで特徴量をぼかすにはMean Poolingを使っていた
しかし、今回のモデルの精度向上のため、meanプールからattention poolingへ変更した
それは、特徴量ごとの重要度を学習しながら集約する事が必要だったから
ちなみに、mean poolingからmulth-head attention poolingに変更したらかなり精度が向上した
そのため、Attention Poolingを理解するためのAttention自体の整理を行った
また、Multi-Head Attention系も必要だったのでそれの理解も載せる

ConvとPool

ConvとPoolの理解

自分の理解だと前提として、普段よく使うConvとPoolは次の理解は以下になる。

Conv
- translation equivariant な性質
- 入力を平行移動すると、特徴マップも対応して平行移動
- 局所特徴を抽出する役割
Pool
- 特徴量を圧縮・集約する役割
- Max Poolingなどは一般に非可逆
- Global Poolingでは位置に対するinvarianceを得られる

Mean Poolingの理解

例えば、Mean Poolingでは、入力特徴量の平均による集約という役割。

$$ z = \frac{1}{N}\sum_{i=1}^{N}x_i $$

WHere:

$x_i$ が各位置の特徴量
$N$ が特徴量の個数
$z$ が集約後の特徴量

つまり、

すべての位置に同じ重み $\frac{1}{N}$
学習されない重み
特徴量の重要度を区別しない集約

例：3個の特徴量を集約する場合

$$ z = \frac{x_1+x_2+x_3}{3} $$

$x_1$ が重要でも $x_3$ が重要でも同じ扱い
「全部を平等に見る」という仮定となる

Attentionの仕組み

Query・Key・Value

Attentionの中心となるのは、Query、Key、Valueで、Q、K、V と表記されるベクトル。

それぞれの意味:

Query
- 「何を探しているか」を表すベクトル
Key
- 「検索対象がどのような特徴を持つか」を表すベクトル
Value
- 実際に取得する情報を表すベクトル

QueryとKeyの比較による関連度の計算、そして、関連度を使ったValueの重み付き和。

なぜQueryとKeyを分けるのか

単純な特徴量同士の類似度だけではなく、「何を探すための表現か」と「検索されるための表現か」の分離をしたいから。

また、同じ入力でも異なる役割を持たせるための線形変換となっている（これは検索も検索対象も同じXから作るのでself-attention）:

$$ Q = XW_Q $$ $$ K = XW_K $$ $$ V = XW_V $$

それぞれ、以下を意味する。

$W_Q$によるQuery用の表現への変換
$W_K$によるKey用の表現への変換
$W_V$によるValue用の表現への変換

つまり、Attentionする情報も、$W_q, W_k, W_v$ですべて学習可能になっているということ。

Attentionの処理

Transformerなどで使われる代表的なAttentionの仕組み（Scaled Dot-Product Attention）。

$$ \mathrm{Attention}(Q,K,V) = \mathrm{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V $$

Where:

$QK^\top$
- Queryと Keyの内積
- Queryと各Keyの相性の計算
- つまり、1つのQueryに注目した場合、Query:Key = 1:Nで計算を行う
- 値が大きいほど高い関連度
- 「このQueryに対して、どのKeyが関係しているか」というスコア
$\sqrt{d_k}$ による除算
- Key の次元数を $d_k$ としたときのスケーリング
- 次元数が大きくなることで内積の値が大きくなりやすい問題への対処
- Softmaxの出力が極端になることを抑えるための処理
Softmax
- スコアから Attention Weight への変換 $$ \alpha_i = \frac{\exp(s_i)} {\sum_j\exp(s_j)} $$
- $s_i$ が Query と各 Key の関連度
- $\alpha_i$ が変換後の Attention Weight
- 重みの合計が1になる性質
- 「どこをどれくらい見るか」を表す割合
Valueとの積
- Attention Weight を使った Value の重み付き和 $$ z = \sum_i\alpha_iV_i $$
- 関連度の高い Value を強く反映した出力
- Query に応じて変化する集約結果

Self-Attention

Self-Attentionとは

最も一般的なAttentionがself-attetion
これはI have a penみたいな文章のattentionに使われている
例えば、haveと 4つ（I, have, a, pen）の関連性を取るため
（ちなみに、この場合は、haveとpenが高いattention weightになると思う）

$$ Q = XW_Q $$ $$ K = XW_K $$ $$ V = XW_V $$

Where:

$X$は元々の行列（i have a penなら、4つのベクトルを持つ行列）

Self Attentionの次元

入力が $N$ 個の特徴量の場合： $$ X \in \mathbb{R}^{N \times D} $$

Self-Attention 後も基本的には $N$ 個の特徴量がある。 $$ N \times D \rightarrow N \times D $$

CNNととself-attentionの違い

特徴を取るのに一般的に使われるのはCNN系のカーネル。

そのCNN系とself-Attentionの違いは以下:

CNN
- 近傍を重視する構造
- その近傍のサイズはKernel sizeによって決まる
- また、inductive biasが強み
Attention
- Attentionは「どこと関係が強いか」を重視する構造
- Attention機構によって重要視する所も学習が可能
- つまり、関係性をデータから学習する柔軟性が強み

CNNのInductive bias

画像なので、以下の前提が、CNNアーキにはあるということ:

近くの画素同士は関係が強いはず、という局所性
同じ特徴検出器を画像全体で使い回すというweight sharing
位置がずれても同じ特徴を検出できるというtranslation equivariance

これは全体をくまなく見るAttentionとは違うbiasがかかっている。

Multi-Head Attention

Multi-Head Attentionとは

Attention を1回だけではなく複数並列に実行する仕組み
わかりやすく言えば、QKVの重みがn個あるということ
それだけ別々の特徴を取れるということ

$$ \mathrm{head}_i = \mathrm{Attention} ( QW_i^Q, KW_i^K, VW_i^V ) $$

特徴:

Head ごとに異なる変換
Head ごとに異なる Attention

各 Head の出力を結合する処理は以下となる。

$$ \mathrm{MultiHead}(Q,K,V) = \mathrm{Concat} ( \mathrm{head}_1, \dots, \mathrm{head}_h )W^O $$

つまり、1つのAttentionだけに限定しない関係性の捉え方。

multi-head attentionのイメージ

例えば画像の場合だと、以下を別々に学習するようなイメージ:

形状に反応
テクスチャに反応
特定領域に反応
離れた領域同士の関係に反応

人がHeadごとに役割を定義するわけではなく、複数の表現空間から同じ入力を見る仕組み。

Attention Pooling

Query Attention Pooling

Attention Poolingのよく使われる実装方法の一つがQuery attention pooling
これはPooling用のQueryを用意する方法
ただし、Self Attentionとは違い、入力XからQueryは作らない
そして、ある1個のQueryから複数の特徴量を参照する構造

入力特徴量:

$$ X = [x_1,x_2,\dots,x_N] $$

Pooling用Query（Xから作っていないのがミソ）:

$$ Q = q_{pool} W_p $$

入力からself-attentionと同じようにKeyとValueを生成する:

$$ K = XW_K $$ $$ V = XW_V $$

QueryとKeyの関連度を計算する:

$$ s = \frac{QK^\top}{\sqrt{d_k}} $$

Attention Weightを算出する:

$$ \alpha = \mathrm{softmax}(s) $$

self=attentionのように、Valueの重み付き和を出す:

$$ z = \alpha V $$

最終的に1個のベクトルを得るというpooling構造:

$$ N \times D \rightarrow 1 \times D $$

Query Attention Poolの特徴

大きな違いは、qがxから生成されないということ
そして、次元が1次元かに圧縮されるということ
これはflatten（$N \times D$）ではなく、圧縮（$1 \times D$）となっている
（ゆえにPoolingということ）

Self AttentionとQuery Attention PoolのQueryの違い

Self-AttentionとQuery Attention Poolingの違いとしては、Queryの数
Self-Attentionでは各入力位置に対応する複数のQuery
- 例えば、I have a penなら、4つのqueryができるはず
他方、Query Attention Poolingでは1個の学習可能なQueryを使う構成
つまり、Pooling用Query自体をモデルのパラメータとして学習する方法

学習可能なQueryの直感的な意味:

「このタスクを解くために、入力からどのような情報を集めたいか」を表すベクトル
学習によって更新される検索条件
入力ごとに異なるKeyと比較される共通のQuery
Query自体は同じでも入力ごとに変わるAttention Weight
Keyが入力から生成されることによる動的なPooling

Multi-Head Attention Pooling

Attention PoolingをMulti-Head化したもの
複数の観点から特徴量を集約する仕組み

ヘッドの数による違い:

1つの Head
- 1つの表現空間でのQueryとKeyの比較
- 1種類のAttention Weight
- 1つの観点からの特徴集約
複数の Head
- 複数の表現空間での比較
- Head ごとに異なる Attention Weight
- 複数の観点から得られる集約結果

Multi-healの意味:

各Head の出力を結合した最終表現
単一の重要度だけでは表現しにくい情報への対応

Mean Pooling とのイメージ上の違い

Mean Pooling：全員の意見を同じ重さで平均
Attention Pooling：必要そうな人の意見を強く反映
Multi-Head Attention Pooling：複数の観点から重要な意見を集めて最後に統合

Attentionの意味

Attention Poolingのコア

Attention Poolingのコアは「どこを見るか」という考え方
つまり、入力を均等に扱うのではなく、重要な部分に大きな重みを与える仕組み
Mean Pooling における固定値 $1/N$ を、入力から計算した重みに置き換えるイメージ

$$ z = \sum_{i=1}^{N}\alpha_i x_i $$

Where:

$\alpha_i$ が各特徴量に対する重みで、Attention Weight と呼ばれる値

多くの場合に満たす条件: $$ \sum_i \alpha_i = 1 $$

Mean PoolingとAttention Poolingの違い:

Mean Pooling も一種の重み付き和
Mean Pooling ではすべての $\alpha_i = 1/N$
Attention Poolingでは入力に応じて変化する $\alpha_i$
「何を見るべきか」をモデルが決める仕組み

これがAttention Poolingのコア。

Mean PoolとAttention Poolの違い

Mean Pooling
- 全部を均等に見る Pooling
Attention Pooling
- 入力に応じて見る場所を変える Pooling

Attentionで実際に学習されるもの

Softmax 自体には基本的に学習パラメータなし
内積自体にも学習パラメータなし
主な学習対象となる Q、K、V の射影

$$ W_Q, ; W_K, ; W_V $$

Multi-Head Attention における Head ごとに異なる射影
Head を結合した後の出力変換

$$ W_O $$

Attention Pooling で学習可能なQueryを使う場合のQuery自体
学習によって変化する関連度を測るための空間
学習によって変化する「何を見るか」の基準
Attention Weight 自体を直接パラメータとして保存する仕組みではない点
入力と学習済みパラメータから毎回計算されるAttention Weight

Lllama上のAttentionの例

Attentionの仕組み

まずは、KとQの内積を求める
内積のスカラーは次元数が大きくなるほど大きな値を取りやすいため、$\sqrt{d_k}$ で除する
scale調整ということ

入力はThinking Machinesで、それぞれembeddingして、$x_1$と$x_2$にしている
$X=[x_1, x_2]^T$ということ
それぞれ、QKVを計算する
- $q_i = x_i W_q$
- $k_i = x_i W_k$
- $v_i = x_i W_v$
その後、Attentionをかける
- $q_1$ を、Thinking自身の $k_1$ だけでなく、Machinesの $k_2$ とも比較する
- （この例の場合は、thinkingの$q_1$を更新したい）
つまり、次を計算している
- $q_1 \cdot k_1$: Thinkingから見たThinkingとの関連度
- $q_1 \cdot k_2$: Thinkingから見たMachinesとの関連度
そして、softmaxをかけて、$[14, 12]$を確率と変換している
つまりThinkingのQueryは以下となったということ
- Thinking自身へのAttention → $0.88$
- MachineへのAttention → $0.12$
これらの$a_i=[0.88, 0.12]$がthinkingのattention（αttention weightになる）
Vは$v_i$として学習済みなので最終的に以下になる
$z_i = a_i V$
つまり，$z_1= 0.88 v_1 + 0.12 v_2$

なお、attention weightの定義は以下となる。

$$ \Alpha = \mathrm{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right) $$

$$ \mathrm{Attention}(Q,K,V) = \Alpha V $$

最後に、$Z_i$をconcatして、それを学習するためのoutput重み$W_o$で行列積をする
結果$Z_o = Z W$を得る（Multi-Head Self-Attention outputの意味）

Attentionのパターン

今回のはMHAだったが、他にもパターンがある。

MHA
- 表現力を多く持たせやすい
- KV Cacheが大きい
GQA
- 複数のQ HeadでK/Vを共有
- 表現力と効率の中間
MQA
- 全Q HeadでK/Vを共有
- KV Cacheを大きく削減

LLamaの実際のコード

https://github.com/meta-llama/llama/blob/689c7f261b9c5514636ecc3c5fefefcbb3e6eed7/llama/model.py#L176

実際は、(batch_size, sequence_length, number_of_heads, head_dimension)$みたいな形で処理されている。

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

        self.wq = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wk = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wv = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wo = RowParallelLinear(
            args.n_heads * self.head_dim,
            args.dim,
            bias=False,
            input_is_parallel=True,
            init_method=lambda x: x,
        )

        self.cache_k = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)

        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

        # repeat k/v heads if n_kv_heads < n_heads
        keys = repeat_kv(
            keys, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)
        values = repeat_kv(
            values, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)

        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
        keys = keys.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        values = values.transpose(
            1, 2
        )  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        return self.wo(output)

その他

Attentionの種類の違い

Attention:

Self-Attention
- QとKとVが同じ入力から作られる
- 例えば、ある文章の中で単語の繋がりなどの関連性が大切などのとき
Cross-Attention
- QueryとKey / Valueの出どころが異なるAttention
- $Q$は入力$X$、$K,V$は別の入力$Y$から作られる
- 例えば文章から画像特徴を参照する場合など
Attention Pooling
- Attention を使って複数の特徴を集約する

AttentionとPoolingのサイズの違い

Attention: $$ N \times D \rightarrow N \times D $$

複数の特徴量から複数の特徴量への変換
特徴量同士の関係を使った各特徴量の更新

Pooling: $$ N \times D \rightarrow 1 \times D $$

複数の特徴量から1個の特徴量への集約
特徴量の個数を減らす処理

Attentionの直感的な理解

Attentionの本質としての検索
Query
- 欲しい情報を表す検索条件
Key
- 検索対象を表す見出し
Query × Key
- 欲しい情報との一致度
Softmax
- 一致度から参照割合への変換
Value
- 実際に取得する中身
Attention
- Queryに合ったValueを集める処理
- QueryとKeyによる「どこを見るか」の決定
- Valueによる「何を取り出すか」の決定
そして、QKVは学習することによって形を変える

まとめ

Mean PoolingからAttention Poolingへの変更は、「情報の集約方法を学習可能にする」ということ
具体的な例がないとイメージはつきにくい