はじめに
ニューラルネットワークのモデル圧縮では、不要な重みを削除するプルーニングがよく利用されています。
代表的な方法の一つが、重みの絶対値を基準に小さい重みを削除するMagnitude-based Pruningです。
今回は、この「削除する前に重みを別の表現に変換する」という方法を試してみます。
具体的には、重みをFFT（高速フーリエ変換）によって周波数領域に変換し、そこで係数を削減した後、逆FFTによって重みに戻します。
そして、通常のMagnitude-based Pruningと比較して、削減後のモデルの精度がどのように変化するかを確認します。
実験には、ResNet18とCIFAR-10を使用しました。
FFTとは？
FFT（Fast Fourier Transform：高速フーリエ変換）は、データを「どのような周波数成分から構成されているか」という形に変換する手法です。
例えば音声や画像などのデータは、そのまま扱うだけでなく、
- 低い周波数の成分
- 高い周波数の成分
といった形に分けて考えることができます。
今回の実験では、ニューラルネットワークの重みをFFTによって周波数領域に変換し、その中から振幅の小さい成分を削除します。
その後、逆FFTによって再び重みの形に戻します。
ざっくり書くと、
ニューラルネットワークの重み
↓
FFT
↓
周波数領域
↓
小さい成分を削除
↓
逆FFT
↓
削減された重み
という流れです。
今回使用したFFTは、PyTorchのtorch.fftを利用しています。
比較する手法
今回は以下の2つを比較しました。
Magnitude-based Pruning
比較対象として、重みの絶対値に基づくMagnitude-based Pruningを使用しました。
これは、
絶対値が小さい重みから削除する
というシンプルな方法です。
例えば、
0.8
0.03
-0.01
0.5
-0.02
という重みがあった場合、絶対値の小さい値を削除します。
FFT Magnitude Pruning
今回試した方法です。
まず重みをFFTして周波数領域に変換します。
その後、
FFT後の係数の振幅が小さいものから削除する
という処理を行います。
最後に逆FFTして、削減された重みを得ます。
つまり、
重み
↓
FFT
↓
周波数成分
↓
振幅の小さい成分を削除
↓
逆FFT
↓
重み
という処理です。
通常のMagnitude-based Pruningでは「重みそのもの」の大きさを基準にします。
一方、今回の方法では「FFTによって得られた係数」の大きさを基準にします。
この違いによって、同じ割合まで削減した場合に性能差が生まれるのかを調べます。
実験環境
今回使用した環境は以下です。
- GPU: NVIDIA GeForce RTX 3080
- フレームワーク: PyTorch
- モデル: ResNet18
- データセット: CIFAR-10
- 学習: 20 epochs
- 削減後の再学習: なし
CIFAR-10は10クラスの画像分類データセットで、32×32ピクセルのカラー画像60,000枚から構成されています。
まずResNet18をCIFAR-10で学習させました。
学習後のテスト精度は、
81.15%
でした。
今回は「削減した後にどこまで精度を維持できるか」を比較したかったため、削減後の再学習やFine-tuningは行っていません。
実装
Magnitude-based Pruning
まず通常のMagnitude-based Pruningです。
def magnitude_prune_tensor(tensor, keep_ratio):
if tensor.numel() <= 1:
return tensor
flat = tensor.abs().flatten()
k = max(1, int(flat.numel() * keep_ratio))
threshold = torch.topk(flat, k).values.min()
mask = tensor.abs() >= threshold
return tensor * mask
モデル全体に適用します。
def magnitude_prune_model(model, keep_ratio):
model = copy.deepcopy(model)
with torch.no_grad():
for param in model.parameters():
param.data = magnitude_prune_tensor(
param.data,
keep_ratio
)
return model
FFT Magnitude Pruning
次にFFTを使った方法です。
def fft_compress_tensor(tensor, keep_ratio):
if tensor.numel() <= 1:
return tensor
if tensor.ndim == 1:
freq = torch.fft.fft(tensor)
else:
freq = torch.fft.fft2(tensor)
magnitude = freq.abs().flatten()
k = max(1, int(magnitude.numel() * keep_ratio))
threshold = torch.topk(magnitude, k).values.min()
mask = freq.abs() >= threshold
freq = freq * mask
if tensor.ndim == 1:
restored = torch.fft.ifft(freq).real
else:
restored = torch.fft.ifft2(freq).real
return restored.reshape(tensor.shape)
モデル全体に適用します。
def compress_model_fft(model, keep_ratio):
model = copy.deepcopy(model)
with torch.no_grad():
for param in model.parameters():
param.data = fft_compress_tensor(
param.data,
keep_ratio
)
return model
実験結果
保持する成分の割合を100%から10%まで10%刻みで変化させました。
ここでいう「Keep Ratio」は、
元の重み・FFT係数のうち、何%を残すか
を表します。
例えば50%なら、半分の情報だけを残します。
Keep Ratioが右に行くほど、より多くの成分を削減しています。
結果は以下のようになりました。
| Keep Ratio | FFT Magnitude | Magnitude-based Pruning |
|---|---|---|
| 100% | 81.15% | 81.15% |
| 90% | 81.16% | 81.42% |
| 80% | 80.41% | 80.13% |
| 70% | 79.05% | 76.42% |
| 60% | 78.67% | 65.18% |
| 50% | 71.32% | 48.71% |
| 40% | 57.17% | 17.53% |
| 30% | 34.10% | 10.20% |
| 20% | 12.97% | 10.00% |
| 10% | 10.00% | 10.00% |
グラフを見ると、Keep Ratioが40〜70%付近で差が大きくなっています。
例えば、**Keep Ratio 50%**では、
- FFT Magnitude: 71.32%
- Magnitude-based Pruning: 48.71%
となりました。
同じ50%まで削減しているにもかかわらず、今回の実験ではFFTを使ったほうが約22ポイント高い精度になっています。
結果について考えてみる
今回の結果では、重みの絶対値を基準に削減するよりも、
FFTして周波数領域に変換してから、係数の大きさを基準に削減する
ほうが、今回の条件では精度を維持しやすい結果となりました。
特にKeep Ratio 40%では、
FFT Magnitude 57.17%
Magnitude-based Pruning 17.53%
となり、大きな差が出ました。
一方で、どのKeep RatioでもFFTのほうが優れているわけではありません。
Keep Ratio 90%では、
FFT Magnitude 81.16%
Magnitude-based Pruning 81.42%
となり、通常のMagnitude-based Pruningのほうがわずかに高い精度でした。
また、Keep Ratioが20%以下になると、どちらの手法でもほぼランダム推測レベルまで精度が低下しています。
CIFAR-10は10クラス分類なので、ランダムに予測した場合の精度は約10%です。
「低周波を残せばいい」のか？
FFTを使うのであれば、
「重要なのは低周波成分なのでは？」
という疑問も出てきます。
そこで、FFT係数の大きさではなく、低周波成分だけを残す実験も行いました。
結果は以下のようになりました。
| Keep Ratio | Low Frequency |
|---|---|
| 100% | 58.47% |
| 50% | 29.30% |
| 20% | 10.03% |
| 10% | 10.00% |
単純に低周波成分だけを残した場合、FFT Magnitude Pruningほど良い結果にはなりませんでした。
そのため今回の実験結果だけを見る限り、
「低周波だから重要」
という単純な話ではなく、
FFT後の係数を振幅の大きさで選択することに意味がある可能性
があります。
ただし、これはあくまで今回の実験条件での結果なので、これだけで一般的な性質だと結論づけることはできません。
今回のFFTは簡易的な実装
今回のコードでは、テンソルの次元数に応じてFFTを適用しています。
特にCNNの畳み込み層の重みは、
[out_channels, in_channels, kernel_height, kernel_width]
という4次元テンソルになっています。
今回の実装では torch.fft.fft2() によって最後の2次元に対してFFTを行っているため、厳密には重みテンソル全体を4次元FFTしているわけではありません。
PyTorchのfft2も、デフォルトでは最後の2次元に対して2次元FFTを行う仕様です。
このあたりは、より本格的に検証する場合には改善する必要があります。
まとめ
今回は、ニューラルネットワークの重みをFFTによって周波数領域に変換し、係数の大きさを基準に削減する方法を試しました。
ResNet18とCIFAR-10を用いた実験では、通常のMagnitude-based Pruningと比較して、特にKeep Ratio 40〜70%付近でFFTを用いた方法のほうが高い精度を維持する結果となりました。
一方で、Keep Ratioを下げすぎると、どちらの手法でも精度が大きく低下しました。
今回の実験は一つのモデル・データセット・実装条件で行ったものなので、FFTによる削減が一般的に有効だとまでは言えません。
とはいえ、「重みを別の表現に変換してから削減する」という方法で、通常のプルーニングとは異なる結果が得られたのは興味深いところでした。