始まり
中には正真正銘のエンジニアであるAIエンジニアもいるだろう。一方で、AIエンジニアの中には、深層学習の闇の住人も紛れ込んでいる。
しばしば出現する ReLU だの attention だの意味のわからない非線形処理。Google の特許に恐怖し、Apache だから仕方なく使う TensorFlow。「Transformer はコンボイだろ」というお決まり文句もとっくに新鮮さを失っている。
我々のプロジェクトはあるドメインのAI認識処理の内製化である。前期まで外部の研究室との共同で開発し、外部コンペでも高い成績を収め、満足のまま終わった。
共同と言っても、向こう側が具体的な学習処理を進めていて、こちら側は仕様やらに口を出す程度。経費や権利を管理する人達ならともかく、平のエンジニアから見たら外注といって差し支えない。次期は内製化し、さらに同様にコンペで成績を出せとのお達しだ。
確かに、ここ数年で深層学習の取り扱いは格段に楽になった。レイヤーを組み合わせればモデルができ、学習用の関数を実行すればパラメータが学習される。だが、コンペ上位クラスの処理はそういう次元ではないのだ。
深層学習は要は、異常な数のパラメータを持つ関数に対する、微分による非線形最適化である。
コンペ上位クラスに挑むということは、微分をするという作業に挑むということだ。おそらく、例の研究室の人たちは、そのあたりも真面目に考えて学習して処理したのだろう。報告書には、訳の分からない数の GPU の並列などが記載されている。火災が心配になるレベルである。
そんなこともあってか、上司から内製化の話を聞いた瞬間、脳が凍りついた。言葉が詰まる。
「…………いやー……ちょっとそれは……冒険……かも……ですね」
え、そうなの？ とか、上が言ってるから決まっちゃうと思うよ、とか返答される。
……まぁ、上が言っているなら結果が見積もれなくてもいいのか？
話を聞くと、どうやら PyTorch を使って工夫すればもっと性能が出る、ということらしい。確かにコンペなら Google の恐怖から解放され PyTorch でも問題はないだろう。しかし、PyTorch ならできて TensorFlow ならできないこととは？ 禅問答でもされているのだろうか。
とにかく、ここにまた、深層学習の性能を取り繕う議論バトルの種が誕生したのである。
必然の進捗
あれから、上司に学習結果の進捗を問われる日々。
精度はもっと上がらないのか、モデルのサイズはもっと小さくならないのか。それってトレードオフやん。
AIエンジニアなら超効率の学習とかして性能を上げろよ、とも。
正直、負け戦だとは思う。しかし、こちらも仕事。何もしないわけにはいかない。
CNN から Transformer への移行など、世の中ではたまにAIの大きな進化が起こる。彼らがどのように考えて新しい技術を成功させてきたのか。超効率の学習……平凡エンジニアの拙者程度にも理解できるものなのだろうか。
そんな時に目にするリファレンスからのメッセージ。
Model.fit の処理をカスタマイズする
教師あり学習を実行するときに fit() を使用するとスムーズに学習を進めることができます。
独自のトレーニングループを新規で書く必要がある場合には、GradientTape を使用すると、細かく制御することができます。
独自のトレーニングループ……気になるワードだ。GradientTape ……お前は、拙者を闇から救ってくれるのか？
class CustomModel(keras.Model):
def train_step(self, data):
# Unpack the data. Its structure depends on your model and
# on what you pass to `fit()`.
x, y = data
with tf.GradientTape() as tape:
y_pred = self(x, training=True) # Forward pass
# Compute the loss value
# (the loss function is configured in `compile()`)
loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
# Compute gradients
trainable_vars = self.trainable_variables
gradients = tape.gradient(loss, trainable_vars)
# Update weights
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
# Update metrics (includes the metric that tracks the loss)
self.compiled_metrics.update_state(y, y_pred)
# Return a dict mapping metric names to current value
return {m.name: m.result() for m in self.metrics}
……ふむ。
optimizer の引数に trainable_vars と gradients があるということは、gradients は loss から単純に計算した微分だということか。返り値はないが、self に続いている trainable_variables を更新するのだろう。
metrics は評価指標だから本質的に学習に関係しているわけではない。一旦は別として考えよう。
……ということは、この trainable_variables をいじれば、特定のレイヤーを学習しないとか、逆に重みを２倍にするとか、そういうことができるわけか。
GradientTape のところにある y_pred = self(x, training=True) はAIの推論処理で、こういう書き方ができるのは、事前に登録されているレイヤークラスのおかげでできているのだろう。
これって、自分でやろうと思うと、self.trainable_variables の配列の次元数とか考えてやらないといけないから大変だろう。デバッグももちろん、行列積の微分とか頭がおかしくなりそう。そしてそこまで頑張っても、レイヤーごとの重みを変えるとかそういうのは、よほど正当な理由がない限り行き当たりばったりで大抵意味がないんだよな。
そういえば、Transformer の attention って……。
全然 CNN と構造が違うよな。自然に行列積のときに次元が合うような行数列数の選択とかもされるだろうし、そういうのも自力ですべてクリアして学習の成功につなげないといけない。基本は完璧でないと厳しいだろう。
Transformer の最初の論文 には効果的な構造だとある。CNN の畳み込みと違って、Transformer だと値の混ぜ方が式(6)みたいにパラメータがない Softmax だったりするから、パラメータを抑えてデータ全体を考慮するようになるのだろう。
そもそも論として、CNN の段階でも構造は Figure 5. みたいに複数出てきて、適した構造というのはプロの視点で考えられている。素人が適当に考えていい構造になる確率は、あまり期待しないほうがよく思える。
超効率の学習の前例として、Transformer を考えたやつは CNN が主流の状態で、こんなことに取り組んでいた、ということになる。
attention レイヤーがない当時は、自分で微分方法も考えて、GradientTape を書いて。一体どんな世界線なんだ。人間が存在する世界なのだろうか。「AIエンジニアなら」とか言われてこんな仕事させられたら詰むんだが。
こんなの無理だ。平凡エンジニア（拙者）が精度を上げようと思ったら、レイヤーを増やすか、教師データを増やすか、それくらいしかない。
コンペ目当てならモデルサイズも制限される。教師データを考えるほうが長期的にマシか。こう言うしかないか……。あと Transformer やべぇ、とも報告しておこう。
……これでも学習考えろって言われるよな、きっと。せめて、拙者よりは線形代数をできる人材を出してくれと言うか……。 attention の式を見てちゃんと微分できないと無理だよな。拙者には無理だし。
毎日のように行われるプロジェクト期限やチェックポイントの議論バトル。文言をアメーバのごとく変え、徐々に緩和されていく目標。おれたちのたたかいはこれからだ！
（このプロジェクトはのちに大幅に規模を縮小された）
補足
GradientTape は手を入れたことはありますが、まさに今回の内容のような気分になりました。便利なモデルを使わせてくれる先人への感謝が溢れます。
上層部との技術的な実行可能性の温度感が合わないのは、結構あったりする……のでしょうか。データ不足です。
AIエンジニアの闇はあまりに深く、頑張っても成果が出ないことはあると思います。それでも、自分の成長という視点も合わせもって少しずつ続けていくことが大事だと思います。