高速フーリエ変換(FFT),数論変換(NTT)による畳み込みの高速化

畳み込み積分

関数 $f,g$ の畳み込み積分とは、イメージとしては、「$g$ を平行移動させながら、$f$ と重なった部分についての合成値を、平行移動のずらし幅それぞれについて求めたもの」。

もう少し具体的に、1次元の場合、

連続関数なら $h(x)=\int_{-\infty}^{\infty}f(t)g(x-t)dt$
離散関数なら $\displaystyle h(x)=\sum_{i=0}^{x}f(i)g(x-i)$

わかりやすさのため、1次元の離散的な畳み込みについての例を示す。（実際は連続的な関数でもよい）
離散の場合、関数 $f(i)$ は、数列 $f$ の第 $i$ 項目、と捉えてもよい。

2つの数列 $A,B$ があったとする。前から、$0,1,2,...$ と添字がついているとする。

i  0  1  2  3  4
A  5  2  3  8  1
B  4  7  6  2

この2つを合成して、新たな1つの数列 $C=A \ast B$ を作る。

愚直には、$C=(c_0,c_1,...)$ を0で初期化し、全ての $(a_i, b_j)$ の組 $(i,j)$ について $c_{i+j}$ に $a_i b_j$ を加算していけばよい。

   i   0   1   2   3   4
   A   5   2   3   8   1
   B   4   7   6   2

a0*B  20  35  30  10
a1*B       8  14  12   4
a2*B          12  21  18   6
a3*B              32  56  48  16
a4*B                   4   7   6   2
------------------------------------  ↓各列についての和
   C  20  43  56  75  82  61  22   2

この合成には、様々な活用法がある。

多項式の掛け合わせ
- 2つの多項式 $f(x)=a_0+a_1x^1+a_2x^2+..., \ \ g(x)=b_0+b_1x^1+b_2x^2+...$
- 積 $f(x)g(x) = c_0+c_1x^1+c_2x^2+...$ になる
音響処理
- 例えば $A$ を音声データ、$B$ をリバーブの特性（インパルス）とすると、畳み込むことでリバーブがかかった音声を再現できる
画像認識
- 複数のピクセルを畳み込み1つの値にまとめることで、特徴を効率的に学習できる

変換による高速化と、変換の種類

離散データについて、たとえば $c_0,c_1,...,c_N$ を求めたいとして、全ての $i+j \le N$ なる組について計算していると、$O(N^2)$ かかる。 $N$ が $10^5$ を超えてくると、数秒では終わらない計算量になる。

ここでもし、以下が成り立つような変換 $\mathcal{F}$、およびそれを元に戻す逆変換があれば、

$\mathcal{F}(A \ast B) = \mathcal{F}(A) \circ \mathcal{F}(B)$
ただし、$\circ$ は要素毎の積（アダマール積）

「$A,B$ それぞれを変換」→「要素同士の積を計算」→「逆変換で元に戻す」とすることで、畳み込みできる。

これをおこなうのが、フーリエ変換など。

変換	対象	特徴
フーリエ変換（FT）	連続関数	周波数成分に分解する
離散フーリエ変換（DFT）	離散データ	周波数成分に分解する。高速な変換が存在
数論変換（NTT）	離散データ	mod P 上の値に変換する。高速な変換が存在

特に、離散フーリエ変換には高速なアルゴリズムがあり、高速フーリエ変換（FFT）と呼ばれている。

一方、フーリエ変換はどうしても（入力と出力が整数でも）途中で小数が発生する。
小数が発生するということは、演算誤差が発生するということである。

そこで、一定の条件を満たす素数 $P$ の mod のもと、正確な値を求められるようにしたのが NTT となる。
（一定の条件：$P-1$ が $2$ でいっぱい割り切れる。理由は後述）

変換の直感的な意味合い

離散データの畳み込みを、形式的冪級数を用いて、多項式の掛け合わせとして解釈する。

$f(x)=a_0x^0+a_1x^1+...+a_Nx^N$
$g(x)=b_0x^0+b_1x^1+...+b_Mx^M$

とすると、畳み込み結果は多項式の積となる。

$h(x)=c_0x^0+c_1x^1+...+c_{N+M}x^{N+M}=f(x) \times g(x)$

ところで、$d$ 次多項式 $p(x)$ は、 $d+1$ 個の異なる変数値における値（$p(0),p(1),...,p(d)$ など）が与えられれば、 $x^0,x^1,...,x^{d}$ の係数が一意に定まることが知られている。

また、$x$ に特定の値を代入した結果については、普通に積が適用できる。

$f(3)=5, g(3)=-7$ のとき、$h(3)=-35$

ここから、

$N+M+1$ 以上の $n$ 個の値 $x_0,x_1,...,x_n$ を選ぶ
$f(x_0),f(x_1),...,f(x_n)$ ならびに $g(x_0),g(x_1),...,g(x_n)$ を求める
それぞれかけあわせて $h(x_0),h(x_1),...,h(x_n)$ を求める
復元して、$c_0,c_1,...,c_n$ を求め、$c_0～c_{N+M}$ を採用する

とすることで、畳み込みできることが分かる。

問題は、2.と4.をどうやって高速に行うかというところである。愚直にやれば、行列演算などで次数の3乗などがかかってしまう。

高速な変換方法

$n$ と $x_i$ を以下のように決めると、高速な方法がある。

$n$ は2べき（$2^k$ と表せる数）
$x_i$ は、以下のような性質を持った値 $w$ を使って、$w^i$（$0 \le i \lt n$）と表せる数
- $w^0,w^1,...,w^{n-1}$ は全て異なる値
- $w^n=1$
- 直交性がある
  - つまり、$\displaystyle \sum_{i=0}^{n-1}w^{ki}$ は、$k=0$ のとき $n$、$k \neq 0$ のとき $0$

$n$ が元の数列長 $N,M$ と合っていなくても、適当にゼロ埋めしとけば結果は変わらない。

このような都合のいい数 $w$ は、複素数の世界、mod の世界で、それぞれ「$1$ の $n$ 乗根」として存在する。
複素数を使うのが FFT、mod でやるのが NTT となる。

ただし NTT の場合、「$1$ の $2^m$ 乗根」が存在するためには、法を $P$ として $P-1$ が $2^m$ の倍数でなくてはならない。

$998244353$ の場合、$2^{23} \times 119+1$ のため、$2^{23}=8300608$ の長さまでの畳み込みが行える。
一方、$1000000007=2 \times 50000003 + 1$ のため、法によっては全く使えない。

で、実際にどうやって高速化しているのかは、以下に行列を使って逐次的に丁寧に説明してくれている記事がある。

うさぎでもわかる信号処理・制御工学　第14羽　高速フーリエ変換(FFT) | 工業大学生ももやまのうさぎ塾

変換を行列の積で表現し、偶数行と奇数行をわけると、それぞれに上手い繰り返し構造が見えることを利用している。

実装例

np.fft の利用

Python では、numpy.fft という関数モジュールがある。

「結果が、複素数の型 float64 の精度、$2^{53}$ 以内に収まる」ことがわかっているなら、こちらを使っても正確な値が求められる。

$f,g$ の最大値を $m$、長さを $n$ として、結果の値は $nm^2$ くらいとなる。

Python

import numpy as np

f = np.array([5, 2, 3, 8, 1])
g = np.array([4, 7, 6, 2])
fft_len = 8  # |f|+|g|-1 以上の2の冪数

# f,gをフーリエ変換
Ff = np.fft.rfft(f, fft_len)
print(Ff)
Fg = np.fft.rfft(g, fft_len)
print(Fg)

# 各要素について掛け合わせる
Ffg = Ff * Fg
print(Ffg)

# 逆変換後、整数に丸める
fg = np.fft.irfft(Ffg, fft_len)
fg = np.rint(fg)
print(fg)

result

F(f)
[19.         +0.j         -0.24264069-10.07106781j
  3.         +6.j          8.24264069 -4.07106781j
 -1.         +0.j        ]

F(g)
[19.         +0.j          7.53553391-12.36396103j
 -2.         -5.j          0.46446609 -0.36396103j
  1.         +0.j        ]
  
F(f)*F(g)
[ 361.         +0.j         -126.34671709-72.89087297j
   24.        -27.j            2.34671709 -4.89087297j
   -1.         +0.j        ]

f*g
[20. 43. 56. 75. 82. 61. 22.  2.]  : 上記の計算結果と一致している

正確なmod上の計算

$10^9$ くらいの mod なら、カラツバ法の要領で、上下30bitに分解して3回畳み込めば正確に求めることができる。

[数学・numpy] 高速フーリエ変換(FFT)による畳み込み | maspyのHP

ただし計算量は増えるので、低速なことは受け入れるしかない。

fftとrfftについて

np.fft.fftとnp.fft.rfft - pLog

numpy.fftには、fftとrfftという2種類の関数が存在する。（逆変換もifftとirfftがある）

fftによるフーリエ変換は、$N$ 要素の数列を入れると $N$ 要素の変換結果が返ってくる。

これ、後ろ半分は負の周波数(?)の情報が格納されているらしい。

フーリエ変換は元の数列が実数であれば、正と負の周波数に同じような数字（複素共役）が現れる性質があるため、通常、後ろ半分の情報は無駄である。

この計算を省略するのがrfftで、前半 $\frac{N+1}{2}$ についてのみ計算することで、半分とまではいかないものの、速くなる。

発展

2次元畳み込み

FFT,NTTは2次元でも有効である。

2次元畳み込みとは、多項式の表現で表すと、$f(x,y)=3+2x+y+5xy,g(x,y)=4-3x+y-2xy$ として、$f(x,y)g(x,y)$ を求めたい、というような意味合いとなる。

これは、$f,g$ を2次元配列で表し、

f(x,y)=3+2x+y+5xy  g(x,y)=4-3x+y-2xy
  [[3, 1],           [[ 4,  1],         Arr[i,j] には、x^i y^j の係数が入る
   [2, 5]]            [-3, -2]]

$f,g$ の配列を、結果のサイズ ($(H_1+H_2-1,W_1+W_2-1)$) に拡張する
→行方向にFFTする
↓列方向にFFTする
要素毎に掛け合わせる
↓列方向に逆FFTする
→行方向に逆FFTする

と、求められる。

1次元化

行方向はいいのだが、列方向はメモリ上の位置が飛び飛びになるため、高速化されにくい。

最初に1次元化することで、計算量は変わらないが、実用上、少し高速になる。

$f,g$ の配列を、結果のサイズ $(H_1+H_2-1)\times (W_1+W_2-1)$ の1次元配列に拡張する
- $f[i,j]$ は、$f'[i \times W + j]$ に格納する。ただし $W=W_1+W_2-1$
FFTする
要素毎に掛け合わせる
逆FFTする

こうしても、異なる行の成分が影響して結果が変わる、ということはない。

目次