雑記

雑記

私の誤りは私に帰属します。お気付きの点がありましたらご指摘いただけますと幸いです。
  1. Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, Wancai Zhang. Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting. arXiv preprint arXiv:2012.07436, 2020.

    [2012.07436] Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting
    GitHub - zhouhaoyi/Informer2020: The GitHub repository for the paper "Informer" accepted by AAAI 2021.

f:id:cookie-box:20210304153450p:plain:w60

Transformer は行列 {\rm Softmax}(QK^\top /\sqrt{d}) によって入力の特徴系列を「前後の文脈を考慮した」特徴系列に変えますが、系列長 L が長くなるほどこの行列を用意して適用するのに \mathcal{O}(L^2) の計算量がかかってしまうのですよね？ 文脈を考慮する先が増えるわけですから当然ですが。しかし、それでも何とか計算を間引いて Transformer に長い系列を流したいこともあるわけです。そのような手法としてこれまでどのようなものが提案されてきたのでしょう？ それぞれの手法にどのような得手不得手があるのでしょうか？

Transformer の1ヘッドがやること
  • x \in \mathbb{R}^{L \times d_{\rm in}} を受け取る．
  • Q = x W_Q \in \mathbb{R}^{L \times d}， K = x W_K \in \mathbb{R}^{L \times d}， V = x W_V \in \mathbb{R}^{L \times d_{\rm out}} を計算する．
  • y = {\rm Softmax}(QK^\top /\sqrt{d}) V \in \mathbb{R}^{L \times d_{\rm out}} を計算する（※ Softmax は各行を正規化する）．

f:id:cookie-box:20210304153503p:plain:w60

まず Informer [1] の間引き方を復習しようか。考え方としては、{\rm Softmax}(QK^\top /\sqrt{d}) の各行の中で一様分布から逸脱している u 行を残したいんだね。逸脱度の指標は一様分布とその行との交差エントロピーだね（論文では KL 情報量といっているけど定数項を省いているから結局交差エントロピーなんだよね）。

f:id:cookie-box:20210304153450p:plain:w60

なるほど…一様分布と i 行目の交差エントロピーは、i 行目の確率分布に最適化した情報量の尺度で一様分布を受け取ったときの平均的な情報量…Q の i 行目 q_i と K の j 行目 k_j を縦ベクトルと考えると、

\displaystyle {\rm CE}_i = \frac{1}{L} \sum_{j=1}^L \left[ \log \frac{\sum_{j'=1}^L \exp(q_i \cdot k_{j'} / \sqrt{d})}{\exp(q_i \cdot k_j / \sqrt{d})} \right] = \log \left[ \sum_{j=1}^L \exp \left( \frac{q_i \cdot k_j}{\sqrt{d}} \right) \right] - \frac{1}{L} \sum_{j=1}^L \frac{q_i \cdot k_j}{\sqrt{d}}
ですね。すべての i についてこの {\rm CE}_i を計算してこれが大きい u 行を残せば計算量を節約でき…って、これを計算するのにすべての q_i \cdot k_j が必要ですよね？ これでは計算量 \mathcal{O}(L^2) じゃないですか…。

f:id:cookie-box:20210304153503p:plain:w60

うんだから厳密に {\rm CE}_i を計算するんじゃないんだよね。まず、その式から以下が成り立つことはわかるよね。下限の方は交差エントロピーが最小になるのは i 行目の確率分布が一様分布に一致するときということからわかるし、上限はその行の最大値をつかって抑えただけだし。それでこの最右辺第2項と第3項を \overline{\rm CE}_i とおくみたい。

\displaystyle \log L \leqq {\rm CE}_i \leqq \log L + \max_j \left[ \frac{q_i \cdot k_j}{\sqrt{d}} \right] - \frac{1}{L} \sum_{j=1}^L \frac{q_i \cdot k_j}{\sqrt{d}}
k_j が多変量正規分布にしたがうと仮定すると、 \overline{\rm CE}_i > \overline{\rm CE}_{i'} かつ i 行目の分散の方が i' 行目の分散より大きいなら高確率で {\rm CE}_i > {\rm CE}_{i'} になるんだって（正確には「となるような q_i, q_{i'} の領域が存在する」といっているけど）。

f:id:cookie-box:20210304153450p:plain:w60

…うーんでも、その主張だと結局 {\rm CE}_i と {\rm CE}_{i'} を比較するのに QK^\top の i 行目と i' 行目が完全に必要になりませんか？ それでは意味が―

f:id:cookie-box:20210304153503p:plain:w60

以下の実装をみると QK^\top を完全に計算しているわけではないね。

https://github.com/zhouhaoyi/Informer2020/blob/main/models/attn.py#L47-L68

QK^\top を計算するんじゃなくて、K のうち sample_k 行だけ残した K' で代用している。いや、正確にいうと、Q の各行に対して sample_k 行の選び方を変えているから、L 通りの {K'}_{1}, {K'}_{2}, \cdots, {K'}_{L} を用意して、以下の QK^\top もどきを計算して、これに基づいて \overline{\rm CE}_i が大きい n_top 行を選んでいるね（※ この記事では Q の i 行目を縦ベクトル q_i と考えているから横ベクトルに戻すために転置しているよ）。
\left( \begin{array}{c} q_1^\top {K'}_{1}^\top \\ q_2^\top {K'}_{2}^\top \\ \vdots \\ q_L^\top {K'}_{L}^\top \end{array} \right)
それで、一度 n_top 行を選んだら改めてそれらの行に対しては QK^\top をきちんと計算しているよ。実装を抜き出すと以下だね。

Informer のセルフアテンションの間引き方

f:id:cookie-box:20210419234518p:plain:w680
系列長 8 で間引きを発生させるには factor をデフォルト値の 5 より下げないといけない。factor を 2 にすると n_top = 2 * ceil( log(8) ) = 6 になるから、セルフアテンションの計算対象が一様分布から逸脱した 6 行に絞られるよ。上図では 2, 3 行目の計算が間引かれたけど、この 2 行が本当に一様分布に最も近い 2 行だったかは確認していない。

つづいたらつづく