Kento Nozawa

Kronecker積を使った全結合層のパラメータ削減

博士課程に入って36ヶ月の1ヶ月が過ぎてビビっています。

これを使って重み行列を小さくできると思ったのですが、調べたら既にあったのでブログ成仏です。

導入

Kronecker積(以下kron)という行列演算があります。 Graph x MLだと、JureさんのKronecker graphやその派生系で使われてたりします。そこで知りました。 Wikipediaによるとスパコンのベンチマークの Graph500 でKronecker graphが使われているそうです。

例えば という行列に対するkronは、以下のようになります。

kronで生成される行列は、元の2つの行列の行数、列数をそれぞれとすれば、 行数がで、列数が になります。 このため、kronの演算回数に対して指数的に大きな行列が作れます。 例えば、 に同じ次元をもつ を使って10回kronを行えば、の行列が構成できます。 仮にとすれば、パラメータ4つからこのような大きな行列が構成できることになります。


さて、ある日、kronを使えばニューラルネットのパラメータを減らせるなという気持ちに至ります。 とりあえず全結合層のニューラルネットで実験してみました。


ネットワーク構成

MNISTを3層のMLPで分類します。 構成はこんな感じにしました:

  • Input: 784ユニット
  • Hidden 625ユニット relu
  • Output 10ユニット softmax

重み行列のパラメータ数だけを考えると、

  • Input--Hidden490000
  • Hidden--Output6250

で合計 496250 のパラメータが必要です。 Pytorchで定義するとこんな感じです。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28*28, 625)
        self.fc2 = nn.Linear(625, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

次にkronを使うことを考えます。 まず大きいのは Input--Hiddenです。 という2つ行列にkronを使えば、同じサイズの重み行列が作れます。 このとき、Input--Hiddenの重み行列のパラメータ数はもとが490000に対して2x28x25=1400です。

更に減らしてみましょう。先程の行列を にできます。これでパラメータ数は 700 になりました。 行数と列数は素因数分解できるのでもっと減らせますが、とりあえずこのくらいにします。


次に Hidden--Ouputです。 だったのを にできます。 先ほど違って10のrootが整数ではないので次元数が異なります。 このとき、Input--Hiddenの重み行列のパラメータ数はもとが6250に対して5x25+2x25=175です。

以上のパターンとパラメータ数をまとめます。

名称 パラメータ数 説明
Original ナイーヴなニューラルネット
one-kron-mat Input--Hidden だけkron、分解した行列は同一
two-kron-mat Input--Hidden だけkron
one-kron-mat out-kron 2つのレイヤーでkron、Input--Hiddenで分解した行列は同一
two-kron-mat out-kron 2つのレイヤーでkron

実験

PyTorch examplesのMNISTと同じ設定です。epochは35回としました。

結果

acc loss

  • Input--Hiddenだけをkronで構成しても性能が落ちない
  • Input--Hiddenのkronに使う行列を同一としても性能が落ちない
  • Hidden--Outputをkronで構成するとけっこう負ける

ちなみにkronで構成するパラメータと同じ数のパラメータ数をもつナイーヴなニューラルネットの場合,中間層の次元数は10くらいになり、ここまで性能は出ません。


おわりに

検索したらそれらしい研究が既にありました。

この他のニューラルネットでの使われ方としてnatural gradientの近似に使う方法もICMLやICLRで出ています。


謝辞

議論に付き合っていただいたhimkt氏に感謝です。

Newer >>