はじめての自然言語処理

Transformer のパラメータ数を増やしながらも必要な計算量の増加を抑えることができる Switch Transfomer のご紹介です。Google さんのように1兆6千億パラメータは無理ですが、規模が小さいモデルでも効果が見込めるようなので、実際に動かして確認してみたいと思います。

1. はじめに

今回は今年1月に発表された Switch Transformer ¹ の話です。 Transformer というモデルはそのサイズに応じて性能が伸びることが良く知られています²。近年どんどん巨大化しており、 Switch Transformer ではついにパラメータ数が1兆6千億個に達しました³。

この連載ではこの手の巨大なモデルは「スゴイのはわかるんですけれど、デモをつつくぐらいで手元で動かせないなぁ～。」とスルーしていたのですが、 Switch Transformer はパラメータ数の増加の割に計算量が抑えられる構造になっており、小規模モデルでも効果のあるテクニックだということで試してみることにしました。

ソースも公開されています⁴。ただ公開された指し先は Git のリポジトリではなくて単一の Python ファイル（moe.py）ですね。。。良く見たら、この連載で何回か扱った Mesh Tensorflow ⁵ に含まれるようです。そういえば、 T5 の記事を書くときに Mesh Tensorflow のソースを見ていて「 moe ってなんやろ。。。萌え？」とか思いつつ、その時の本筋と違ったのでスルーしたような、ちょっとだけ中身をみたような？

コレ、どう動かすのかイマイチ確信もてないところもありますが、とりあえずトライしてみます。

まずは Switch Transformer がどんなものなのか、概要のご紹介から始めていきましょう。

2. Switch Transformer

Switch Transformer の根幹となる考え方は「Transformer のパラメータ数をシンプルかつ効率的な計算で最大化する」というものです。

Transformer はサイズに応じて性能が伸びるもので、前述の OpenAI の論文²では、「モデルサイズ」、「データセットサイズ」、「計算量」に関してべき乗則が発見されたとしています。 Switch Transformer の論文¹は「入力サンプルを処理する計算量を固定しつつ、パラメータ数を増加させる」という4つ目の軸を検証するものになります。

この「入力サンプルを処理する計算量を固定しつつ、パラメータ数を増加させる」を実現するために、Switch Transformer は従来の Transformer に Mixture of Experts(MoE) ⁶ の考え方を取り入れています。まずは MoE について簡単に説明します。

Mixture of Experts(MoE)

典型的な深層学習モデルは全ての入力に対して同じパラメータを使用するのですが、 MoE では個別の入力に対して異なるパラメータを用いて計算を行います。 この「入力によって使用するパラメータを切り替える」という動きによって、パラメータ数を大幅に増やしつつ計算量を抑制することを実現しています。

MoE の基本的な計算は以下のとおりです。

まず、x がトークン、W_r がルーティングの重みとして、router の logit は以下のようになります。

これを N 個の expert で softmax して i 番目の expert のゲート値とします。

オリジナルの MoE ではここでゲート値の top-k を用いた加重和として出力を計算します。ここで τ は top-k で選択された expert のインデックス集合です。

このようにオリジナルの MoE では top-k の expert を用いていたのですが、 Switch Transformer ではこれをさらに単純化して top-1 のみ使用しており、これを “Switch Layer” と呼んでいます。

そして、この Switch Layer を通常の Transformer ブロックにおける Feed Forward 層を置き換える形で組み込んでいます⁷。 下図の青い角丸四角形の部分ですね。 FFN 1 ～ FFN 4 が各 expert に相当します。

では、Switch Layer について、もう少し詳しく見ていきましょう。

Switch Layer

Switch Layer で入力シーケンス中のトークンが単一 expert へルーティングされるイメージは以下のようになります。

ここで expert capacity と capacity factor が出てきたので補足しておきましょう。

まず expert capacity ですが、これは文字通り各 expert がトークンを受け取れる容量です。大きく設定すると無駄が大きい（上図右の白セル）ですし、小さくし過ぎると定員オーバではじかれるトークン（上図左の青セルと赤破線）がでてきます。top-1 のみ対象とすることでルーティングの計算や実装を単純化し、expert capacity をより小さく取ることができたとのことです。

それと、「満員電車に乗りそこなった人（上図の赤破線のついた青トークン）はどうなるの？」と心配になりますよね。この場合は 単純に Switch Layer での計算をスキップして、残差接続を通って素通しで次の層に送られることになります⁸ 。

capacity factor はバッチに含まれるトークン数に対する expert capacity × expert 数 の比率ですね。 capacity factor = 1.5 であれば 1.0 はトークンが詰まっていて、0.5 は空き（上記の白セル）です。次は、この空き(=無駄)を抑える為の工夫を見ていきましょう。

ロードバランシングロス

計算効率を上げる為、ルーティングはバッチのトークンを出来るだけ均等に各 expert に配分して欲しいです。そこで Switch Transformer では、以下のようなロードバランシング損失を導入して、元の損失に加算しています（αはハイパーパラメータ、N は expert 数、B はバッチに含まれるトークンの集合、T はバッチのトークン数です）。

f_i は expert i に振られるトークン数の割合です。

P_i はトークンが expert i に振られる確率です。

上記の loss はトークンが各エキスパートに均等に配分されるとき(f_i, P_i が 1/N)に最小になります。論文では α = 10^-2 としています。N が掛かっているのはエキスパート数を変えた場合も loss の値が変わらないようにする為です。f_iは argmax が入っていて微分出来ませんが、P_iの方が微分可能なので loss も微分可能になります。

他にもいろいろ工夫しているようなので、それらも確認しておきましょう。

その他の工夫

Switch Transformer は Switch Layer に加え、転送・計算コスト、学習時の不安定さを抑える以下のような工夫を取り入れています。

選択的精度変更

MoE の適用は学習が不安定になるようで、計算効率の良いデータ型である bloat16 使うとそれがさらに顕著になるようです。そこで、Switch Transformer では ルーティング関数の入り口で float32 へ変換、出口で bfloat16 に再変換することで安定性を確保しています。

より小さい値でのパラメータの初期化

T5 ではパラメータの初期化を平均 μ = 0, 標準偏差 σ = √(s/n) (s はハイパーパラメータでデフォルト値= 1.0 , n は入力のユニット数)の切断正規分布で行っていました。 Switch Transformer では s をデフォルト値の1/10(= 0.1) とすることで、安定性が増し精度も向上することが確認されています。

ファインチューニング時の正則化

Switch Transformer は同等の計算量となる Transformer と比較して非常に多くのパラメータを持つ為、ファインチューニング時の過学習傾向が強まります。これを抑制するために、 expert の計算のみドロップアウト率を 0.4, それ以外を 0.1 としています。

ですが、論文の Table 5. にファインチューニングの実験で使用したモデルの構成が示されています。

Switch-Base のパラメータ数は 7.4B なので論文の記述や次項の図を見るとおそらく 128 experts 構成での実験です。 8 とか 16 とか控えめの expert 数でファインチューニングするときは、 0.4 より抑えた設定にした方が良いかもしれないと思いました。

T5 との性能比較

論文によると T5 Base を基準に expert 数を変更して性能比較した結果は以下のとおりです。

Base サイズでもこれだけ性能上がるなら、試してみようかなっていう気分になっちゃいますよね。

論文には Mesh Transformer を駆使した複数コアへの分散配置の話とか、expert を並列化してパラメータ数が巨大化したモデルを蒸留しても、向上した性能の 30 % は維持できるとか、いろいろ面白いことが書いてあるので、興味のある方は読んでみて下さい。

では、実際に動かしてみましょう。

3. Switch Transformer の実行

さて、実際に動かしてみるわけですが冒頭に書いたとおり、ソースコード的には moe.py がポツンとあるだけで動かし方的な情報は見当たりませんね。。。 どうしたものかと思ったのですが、よく考えてみると、

• moe.py は Mesh Tensorflow で実装されている（つまり gin で設定ができる）
• Switch Transformer は Transformer ブロックの Feed forward を Switch Layer で置き換えたもの
• T5 と同様の事前学習を行い、比較実験している（同じ事前学習のコードを 2回作らないでしょう)

ということは、 t5 のコードをそのまま使って gin の設定で差し替えればよいのではないかという気がしてきました。 また、論文中では事前学習の品質を negative log perplexity で比較しています。そのコードが t5 の中にあったんだっけ?と思い探してみると。。。 ありました。perplexity_eval.gin ⁹です。これでだいたい材料がそろった気がします。

それでは環境を作っていきましょう。 いつものように、記事内のコードスニペットは、特に断りがない場合は Google Colaboratory (以下、Colab)で動かす想定にしています。ノートブックを開き、アクセラレータは TPU を選んで下さい。

Tensorflow は 2.x 系を使います。

%tensorflow_version 2.x 

今回のコードは以下のバージョンで動かしています。 tensorflow を 2.4.1 に差し替えているのは、記事のネタを作っている途中で Colab の tensorflow のバージョンが変わってしまった為です。その時に何やらエラーがでたので、安易にバージョンを戻しました。

!pip install  tensorflow==2.4.1 tensorflow-text==2.4.3 t5[gcp]==0.8.1

この時点で mesh-tensorflow は 0.1.19 が入るのですが、このバージョンの moe.py には軽微なバグがあり動きません。 少し手前のバージョンに差し替えて対応しました。

!pip uninstall -y mesh-tensorflow
!pip install git+git://github.com/tensorflow/mesh.git@9625f34e00a201775249ddb887529da859aa83a8

tensorflow を入れ替えたので、ここまで終わったら一度ランタイムを再起動してください。

チェックポイントを書き出すので GCS の認証を通しておきましょう。

from google.colab import auth
auth.authenticate_user()

事前学習のコーパスとしては日本語 Wikipedia を使いました。 データの品質を考えると Wiki40-B を使ったほうがよいかと思うのですが ¹⁰、第7回 で T5 の検証をしたときと条件を合わせたかったので、こうしています。 Sentencepiece のモデルの作り方も第7回の記事を参考にして下さい。

from t5.data import utils
utils.DEFAULT_SPM_PATH = "gs://somewhere/t5/sentencepiece/wikipedia_20190301_ja_v003.model"
import functools
from t5.data import preprocessors
from t5.data.dataset_providers import TaskRegistry
from t5.data.dataset_providers import MixtureRegistry
from t5.data.dataset_providers import TfdsTask
from t5.data.tasks import DEFAULT_OUTPUT_FEATURES
task_name_wikipedia_ja = "wikipedia_20190301.ja_v003_unsupervised"
TaskRegistry.add(
    task_name_wikipedia_ja,
    TfdsTask,
    tfds_name="wikipedia/20190301.ja:1.0.0",
    text_preprocessor=functools.partial(
        preprocessors.rekey, key_map={"inputs": None, "targets": "text"}),
    token_preprocessor=preprocessors.unsupervised,
    output_features=DEFAULT_OUTPUT_FEATURES,
    metric_fns=[])
MixtureRegistry.add(task_name_wikipedia_ja, [(task_name_wikipedia_ja, 1.0)])

Colab 環境の TPU アドレスを確認します。

import os
import pprint
import json
import tensorflow.compat.v1 as tf

assert 'COLAB_TPU_ADDR' in os.environ, 'ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!'
TPU_ADDRESS = 'grpc://' + os.environ['COLAB_TPU_ADDR']
print('TPU address is', TPU_ADDRESS)
# TPU address is grpc://10.113.179.10:8470

あとは、 gin の設定を書いて事前学習を回すだけです。それっぽい設定が text-to-text-transfer-transformer/t5/models/gin/models/ ¹¹あたりにあるかと思ったのですが、ありませんね。自前で書いていきましょう。

また T5 には 1.1 などの亜種がありますが手元の T5 1.0 相当モデルと比較したいので、今回は T5 1.0 の Feed forward を Switch Layer に書き換えたモデル としました。

%%bash
cat << EOF > bi_base_moe.gin
# from bi_bert_base.gin
utils.run.model_type = "bitransformer"
d_model = 768
num_layers = 12
d_ff = 3072 # ①
num_heads = 12
d_kv = 64

# from bi_v1.gin
transformer.Unitransformer.positional_embedding = False
transformer_layers.SelfAttention.relative_attention_type = "bias_shared"

# my own
import mesh_tensorflow.transformer.moe

encoder/transformer.make_layer_stack.layers = [
    @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
    @mesh_tensorflow.transformer.moe.MoE1D, # ②
]
decoder/transformer.make_layer_stack.layers = [
    @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
    @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
    @mesh_tensorflow.transformer.moe.MoE1D, # ②
]

MoE1D.moe_gating = "switch" # ③
MoE1D.num_experts = 8 # ④
MoE1D.hidden_size = 3072 # ⑤
MoE1D.dropout_rate = 0.1 # ⑥
MoE1D.capacity_factor_train = 1.25 # ⑦

import mesh_tensorflow.layers
VarianceScalingInitializer.scale = 0.1 # ⑧
EOF

ポイントになるところを順に見ていきましょう。

• ① Feed forward 層を差し替えるので d_ff は不要な気がしてますが、意図していないところで参照しているかもしれないので残しています。
• ② Feed forward 層(DenseReluDense) を MoE1D で差し替えます。MoE1D が Mixture of Experts の実装になっており、通常の MoE や Switch Layer, No-Token-Left-Behind⁸ の振る舞いを設定で実現できるコードになっています。MoE2D もありますが、こちらは Self Attention に MoE を適用する場合に使うようです⁷。
• ③ 今回は Switch Transformer の検証をするので switch を選びます。
• ④ ここは expert 数ですね。とりあえず 8 にしました。
• ⑤ ここは各 expert のユニット数です。論文には明確な記載がありませんでしたが、「通常の T5 の Feed forward を N 並行にしたら」みたいな比較がしたいと思ったので、 3072 にしています。
• ⑥ Switch Layer に適用される dropout です。ファインチューニング時はここを 0.4 に設定するという記述がありましたが、T5 1.0 相当の事前学習なので 0.1 にしています。
• ⑦ Capacity factor の設定です。出来るだけ小さくしたいですが、論文の Table 1. を見たところ 1.25 あたりが無難そうだと思いました。
• ⑧ 「その他の工夫」の項で書いたパラメータの初期化を小さくする件です。ソースを見たところ、ここ¹²の設定でいいんじゃないかなーと思いました。

ホントにこれでいいんだろうか？と確信持てていないのですが、とりあえず動かしてみましょう。

事前学習の実行

それでは、事前学習を回していきます（例によって Colab のランタイムの寿命では終わりきらないので、動作確認できたら GCP でインスタンス立てて Cloud TPU で回してください）。

ある程度 expert 数を変えて実験したく、32 experts とかで実験して OOM とか言われるのが怖かったので、 batch_size は控えめな数字（といっても過去の連載時と同じですが）にしています。

from t5.models.mesh_transformer_main import FLAGS
from t5.models.mesh_transformer_main import main
FLAGS.mark_as_parsed()
FLAGS.tpu = TPU_ADDRESS

FLAGS.model_dir = 'gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1'

tf.flags.FLAGS.gin_file=[
  "dataset.gin", 
  "./bi_base_moe.gin",
  "objectives/span_3_15_u_u.gin",
  "learning_rate_schedules/rsqrt_no_ramp_down.gin"]
tf.flags.FLAGS.gin_param=[
  "utils.tpu_mesh_shape.model_parallelism = 1",
  "utils.tpu_mesh_shape.tpu_topology = 'v2-8'",
  "run.batch_size = ('tokens_per_batch', 65536)",
  "run.train_steps = 524288",
  "MIXTURE_NAME = 'wikipedia_20190301.ja_v003_unsupervised'"
]

tf.disable_v2_behavior()
tf.logging.set_verbosity(tf.logging.INFO)
main([])

# WARNING:tensorflow:From /usr/local/lib/python3.7/dist-packages/tensorflow/python/compat/v2_compat.py:96: disable_resource_variables (from tensorflow.python.ops.variable_scope) is deprecated and will be removed in a future version.
# Instructions for updating:
# non-resource variables are not supported in the long term
# WARNING:tensorflow:From /usr/local/lib/python3.7/dist-packages/tensorflow/python/compat/v2_compat.py:96: disable_resource_variables (from tensorflow.python.ops.variable_scope) is deprecated and will be removed in a future version.
# Instructions for updating:
# non-resource variables are not supported in the long term
# ERROR:root:Path not found: dataset.gin
# ERROR:root:Path not found: objectives/span_3_15_u_u.gin
# ERROR:root:Path not found: objectives/span.gin
# ERROR:root:Path not found: objectives/denoise.gin
# ERROR:root:Path not found: learning_rate_schedules/rsqrt_no_ramp_down.gin
# WARNING:absl:wikipedia_20190301.ja_v003_unsupervised is both a Task and a Mixture, returning Mixture
# ...

動かすと挨拶代わりに、ERROR:root:Path not found: dataset.gin と絶対ダメそうなメッセージが出力されて「あぅっ」となりますが、出力された operative_config.gin などを見ていると大丈夫そうです。

あと、この設定では出力された checkpoint が消されずにどんどん蓄積していきます。 expert 数を大きく取ったときは GCS の容量が結構なことになるので、気を付けてくださいね（negative log perplexity の曲線を書こうと思うと消されても困るんですけど）。

では、事前学習が終わった体で精度評価しましょう。

精度の評価

精度評価は前述の perplexity_eval.gin を使えば良さそうです。 ノートブックを開き、アクセラレータは TPU を選んで TPU アドレスを確認するところまでは同じです。

以下のようにして精度評価を実行します。こちらは Colab で動かしても大丈夫な時間で完了したかと思います。

from google.colab import auth
auth.authenticate_user()

from t5.models.mesh_transformer_main import FLAGS
from t5.models.mesh_transformer_main import main

FLAGS.mark_as_parsed()

FLAGS.tpu = TPU_ADDRESS
FLAGS.model_dir = 'gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1'
tf.flags.FLAGS.gin_file=[
  os.path.join(FLAGS.model_dir, "operative_config.gin"),                       
  "perplexity_eval.gin",
  ]
tf.flags.FLAGS.gin_param=[
  "utils.tpu_mesh_shape.model_parallelism = 1",
  "utils.tpu_mesh_shape.tpu_topology = 'v2-8'",
  "run.batch_size = ('tokens_per_batch', 65536)",
  "MIXTURE_NAME = 'wikipedia_20190301.ja_v003_unsupervised'"
]

tf.disable_v2_behavior()
tf.logging.set_verbosity(tf.logging.WARN)
main([])

今回の構成ですと、評価の結果は model_dir 直下の eval_wikipedia_20190301.ja_v003_unsupervised_train_12800 というフォルダに出力されます。

!gsutil ls gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/eval_wikipedia_20190301.ja_v003_unsupervised_train_12800/

# gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/eval_wikipedia_20190301.ja_v003_unsupervised_train_12800/
# gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/eval_wikipedia_20190301.ja_v003_unsupervised_train_12800/events.out.tfevents.1617678544.9d1d5a3f712e 

では、評価結果を確認していきましょう。

4. 評価結果の確認

筆者の環境では出力されたファイルを Tensorboard で読み込んだところ、なにも表示されませんでした。 仕方ないので以下のようにして読み込んでみると、ちゃんと欲しい情報は書き込まれているようです。

import tensorflow.compat.v1 as tf

path_to_events_file = "gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/eval_wikipedia_20190301.ja_v003_unsupervised_train_12800/events.out.tfevents.1617678544.9d1d5a3f712e"

results = []
for e in tf.train.summary_iterator(path_to_events_file):
  for v in e.summary.value:
    if v.tag == "neg_log_perplexity":
      results.append((e.step, v.simple_value))
steps = []
nlppls  = []
for step, nlppl in results:
  steps.append(step)
  nlppls.append(nlppl)

print(steps[:3])
print(nlppls[:3])
# [5100, 10200, 15300]
# [-3.1070008277893066, -2.666059732437134, -2.433549404144287]

expert 数を変えたりしながら何パターンか事前学習を行って精度評価した結果をまとめたものが以下のプロットです¹³。

T5_Base は第7回で作ったモデル、SW_Base(*ex) は基本的には前述の bi_base_moe.gin の設定で expert 数だけ変えたものです。 確かに expert 数を増やすと性能は伸びていますが、16, 32 experts は論文の図に比べると伸びが鈍いですね。。。batch_size が小さいことやデータセットの規模などの影響があったかもしれません。

パラメータ数と実行速度

ここまでの実験からパラメータ数(“Trainable Variables”)と実行速度を学習時のログから拾ってプロットしてみました(下図の青点）。

パラメータ数が増加している割に、実行速度の低下が抑えられているのが確認できますね。ただ 16 experts で 1925438080 パラメータですと、T5 との性能比較で引用した論文 Figure 1 の左図と数が合いません。

そこで、試しに expert のユニット数を半分の 1536 にした 16 experts で動かしてみました（上図の赤点）。正確なところは分かりませんが、Figure 1 の数値とほぼ合っています。確信はないですが expert のユニット数は 1536 で実験したのかもしれないですね。

こうなってくると、「大きい expert が少数」と「小さい expert が多数」でどちらが良いのか気になってきます。さすがに気力がなくなってきたので、最初の100000 ステップだけ動かして比較してみました。

確認できる範囲では、ほぼ性能の優劣はなさそうですね。実行速度としては expert のユニットサイズが小さい方が有利なので 1536 ユニットの 16 expert が良いでしょう。

さて、せっかく事前学習したのでファインチューニングの手順も紹介しておきましょう。

5. ファインチューニングと検証

ここからはファインチューニングをしていきます。

とは言っても、手順は第7回で T5 を動かした時と基本的に同じです。ですが、ライブラリの更新等があって過去のコードそのままでは動かなくなっているので改めて紹介しますね。

シーケンス長やバッチサイズにもよりますが、今回実験したユニット数 3072 の 8 experts くらいのサイズであれば、 GPU 1コアでも動作しました¹⁴。 それでは新たに Colab のノートブックを開き、アクセラレータは GPU を選んで下さい。

セットアップ

GCS の認証を通して Tensorflow は 2.x 系を選んでおきます¹⁵。

from google.colab import auth
auth.authenticate_user()
%tensorflow_version 2.x 

ただ、前章で使用した t5 と mesh-tensorflow の組み合わせはファインチューニングの検証で使用する eval.gin が正常に動作しませんでした。 ですので、t5 を 0.9.1 に更新、mesh-tensorflow も新しいもので動かします¹⁶。

!pip install tensorflow-text==2.5.0  t5[gcp]==0.9.1 git+git://github.com/tensorflow/mesh.git@a54f5cf75ef44d8a97 sacrebleu[ja]==1.5.0

学習データの準備

ファインチューニング用のデータはこの連載で何回か使った、やさしい日本語データセットです。 加工の仕方は第14回、第7回を参考にして下さい。

!gsutil cp gs://somewhere/snow*.tsv .
!wc -l *.tsv
#   7040 snow_t15_23_dev.tsv
#   7040 snow_t15_23_test.tsv
#  56317 snow_t15_23_train.tsv
#  70397 total

やさしい日本語変換タスクの定義とファインチューニングの実行

続いてファインチューニング用のタスクを定義します。

%%bash
cat <<EOF > add_snow.py
import functools
import tensorflow as tf
from t5.evaluation import metrics
from t5.data import preprocessors
from seqio import vocabularies
from t5.data.utils import DEFAULT_EXTRA_IDS
from seqio import Feature
from t5.data.dataset_providers import TaskRegistry
from t5.data.dataset_providers import TextLineTask
from sacrebleu import corpus_bleu

def bleu(targets, predictions):
  predictions = [tf.compat.as_text(x) for x in predictions]
  if isinstance(targets[0], list):
    targets = [[tf.compat.as_text(x) for x in target] for target in targets]
  else:
    targets = [tf.compat.as_text(x) for x in targets]
    targets = [targets]

  bleu_score = corpus_bleu(predictions, targets,
                                     smooth_method="exp",
                                     smooth_value=0.0,
                                     force=False,
                                     lowercase=False,
                                     tokenize="ja-mecab",
                                     use_effective_order=False)
  return {"bleu": bleu_score.score}

task_name = "snow_t15_23"

tsv_path = {
    "train": "./snow_t15_23_train.tsv",
    "validation": "./snow_t15_23_dev.tsv",
    "test": "./snow_t15_23_test.tsv",
}

TaskRegistry.add(
    task_name,
    TextLineTask,
    split_to_filepattern=tsv_path,
    text_preprocessor=[
      functools.partial(
          preprocessors.parse_tsv,
          field_names=["inputs", "targets"]),
    ],
    output_features = Feature(vocabularies.SentencePieceVocabulary(
      "gs://somewhere/t5/sentencepiece/wikipedia_20190301_ja_v003.model",
      DEFAULT_EXTRA_IDS)),
    metric_fns=[bleu])
EOF

ファインチューニング用のディレクトリに事前学習済みモデルをコピーします。

!gsutil cp gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/checkpoint \
  gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/snow/
!gsutil cp gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/ope* \
  gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/snow/
!gsutil cp gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/model.ckpt-524288* \
  gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/snow/

これで準備ができたので、以下のようにしてファインチューニングを実行します。

パラメータ数が増えてるので T5 の時よりも学習量を増やしています。

!export PYTHONPATH=${PYTHONPATH}:. && \
  \
  OPERATIVE_CONFIG='gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/snow/operative_config.gin' && \
  FINE_TUNED_MODEL_DIR='gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/snow' && \
  FINE_TUNING_BATCH_SIZE=`expr 40 \* 8` && \
  PRE_TRAINGING_STEPS=524288 && \
  FINE_TUNING_STEPS=`expr $PRE_TRAINGING_STEPS + 16000` && \
  INPUT_SEQ_LEN=40 &&\
  TARGET_SEQ_LEN=40 &&\
  \
  echo "OPERATIVE_CONFIG=$OPERATIVE_CONFIG" &&\
  echo "FINE_TUNED_MODEL_DIR=$FINE_TUNED_MODEL_DIR" &&\
  echo "FINE_TUNING_BATCH_SIZE=$FINE_TUNING_BATCH_SIZE" &&\
  echo "PRE_TRAINGING_STEPS=$PRE_TRAINGING_STEPS" &&\
  echo "FINE_TUNING_STEPS=$FINE_TUNING_STEPS" && \
  echo "INPUT_SEQ_LEN=$INPUT_SEQ_LEN" && \
  echo "TARGET_SEQ_LEN=$TARGET_SEQ_LEN" && \
  \
  t5_mesh_transformer \
  --model_dir="$FINE_TUNED_MODEL_DIR" \
  --module_import="add_snow" \
  --gin_file="dataset.gin" \
  --gin_file="$OPERATIVE_CONFIG" \
  --gin_param="run.layout_rules=''" \
  --gin_param="run.mesh_shape=''" \
  --gin_param="utils.get_variable_dtype.activation_dtype='float32'" \
  --gin_param="MIXTURE_NAME = 'snow_t15_23'" \
  --gin_file="learning_rate_schedules/constant_0_001.gin" \
  --gin_param="run.train_steps=$FINE_TUNING_STEPS" \
  --gin_param="run.sequence_length = {'inputs': $INPUT_SEQ_LEN, 'targets': $TARGET_SEQ_LEN}" \
  --gin_param="run.save_checkpoints_steps=3200" \
  --gin_param='dropout_rate=0.1' \
  --gin_param='mesh_tensorflow.transformer.moe.MoE1D.dropout_rate=0.4' \
  --gin_param="run.batch_size=('tokens_per_batch', $FINE_TUNING_BATCH_SIZE)" 

T5 の時と異なっているのは、mesh_tensorflow.transformer.moe.MoE1D.dropout_rate=0.4 ですね。 これが expert に適用するドロップアウト率の設定になります。

ファインチューニング済みのモデルができたので、検証を行います。

検証の実行

検証の実行方法は T5 の時とおなじですね。

!export PYTHONPATH=${PYTHONPATH}:. && \
  \
  FINE_TUNED_MODEL_DIR='gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/snow' && \
  OPERATIVE_CONFIG=$FINE_TUNED_MODEL_DIR'/operative_config.gin' && \
  \
  echo "OPERATIVE_CONFIG=$OPERATIVE_CONFIG" &&\
  echo "FINE_TUNED_MODEL_DIR=$FINE_TUNED_MODEL_DIR" &&\
  \
  t5_mesh_transformer \
  --model_dir="$FINE_TUNED_MODEL_DIR" \
  --module_import="add_snow" \
  --gin_file="$OPERATIVE_CONFIG" \
  --gin_param="run.layout_rules=''" \
  --gin_param="run.mesh_shape=''" \
  --gin_file="eval.gin" \
  --gin_file="beam_search.gin" \
  --gin_param="utils.get_variable_dtype.slice_dtype='float32'" \
  --gin_param="utils.get_variable_dtype.activation_dtype='float32'" \
  --gin_param="MIXTURE_NAME = 'snow_t15_23'" \
  --gin_param="run.dataset_split='validation'" \
  --gin_param="run.batch_size=('tokens_per_batch', 320)" \
  --gin_param="eval_checkpoint_step ='all'" 2>&1 | tee eval.log

結果を確認しましょう¹⁷。

!cat eval.log | grep -e "^INFO.*bleu"
# INFO:tensorflow:eval/snow_t15_23/bleu at step 524288: 5.187
# INFO:tensorflow:eval/snow_t15_23/bleu at step 527488: 78.016
# INFO:tensorflow:eval/snow_t15_23/bleu at step 530688: 78.790
# INFO:tensorflow:eval/snow_t15_23/bleu at step 533888: 78.720
# INFO:tensorflow:eval/snow_t15_23/bleu at step 537088: 78.661
# INFO:tensorflow:eval/snow_t15_23/bleu at step 540288: 78.344

!cat eval.log | grep -e "^INFO.*bleu" | sort -k 5,5 -n -r | head -1
# INFO:tensorflow:eval/snow_t15_23/bleu at step 530688: 78.790

最良のチェックポイントを特定したのでテストセットでスコアを確認します。 run.dataset_split で test を指定して eval_checkpoint_step で対象のチェックポイントを指定している以外は先程と同じです。

!export PYTHONPATH=${PYTHONPATH}:. && \
  \
  FINE_TUNED_MODEL_DIR='gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/snow' && \
  OPERATIVE_CONFIG=$FINE_TUNED_MODEL_DIR'/operative_config.gin' && \
  \
  echo "OPERATIVE_CONFIG=$OPERATIVE_CONFIG" &&\
  echo "FINE_TUNED_MODEL_DIR=$FINE_TUNED_MODEL_DIR" &&\
  \
  t5_mesh_transformer \
  --model_dir="$FINE_TUNED_MODEL_DIR" \
  --module_import="add_snow" \
  --gin_file="$OPERATIVE_CONFIG" \
  --gin_param="run.layout_rules=''" \
  --gin_param="run.mesh_shape=''" \
  --gin_file="eval.gin" \
  --gin_file="beam_search.gin" \
  --gin_param="utils.get_variable_dtype.slice_dtype='float32'" \
  --gin_param="utils.get_variable_dtype.activation_dtype='float32'" \
  --gin_param="MIXTURE_NAME = 'snow_t15_23'" \
  --gin_param="run.dataset_split='test'" \
  --gin_param="run.batch_size=('tokens_per_batch', 320)" \
  --gin_param="eval_checkpoint_step = 530688" 2>&1 | tee test.log

最終的な BLEU スコアは以下のとおりです。

!cat test.log | grep -e "^INFO.*bleu"
# INFO:tensorflow:eval/snow_t15_23/bleu at step 530688: 78.982

MoE1D.dropout_rate の値を変更しながら試してみました。そこそこ時間かかるので、各値で1回ずつの試行です。 今回は 0.1 が一番良い結果が得られましたが、真ん中（0.2）が凹んでいるのがちょっと微妙ですね。各設定で複数回試行しつつ、もっと多くのタスクで試さないとなんとも言えないところかもしれません。

第7回で T5 を試したときは 78.80 でした（このときは、 sacrebleu 日本語で使うためにちょっと弄ってますが）。 1 pt 程良くなりはしましたが、モデルのサイズアップや実行速度まで考えると、あまり旨味のない結果になってしまいましたね。 事前学習の精度評価の結果からするともうちょっと効果あって欲しかったので残念です。。。

ただ、論文の Table 6. に T5 と Switch Transformer の下流タスクでの性能比較が出ています。

この表を見ると GLUE で +2.4, SuperGLUE で +4.4 です。この Switch Transformer はパラメータ数≒ 7.4B ですので、それを考えると今回（パラメータ数≒1B）の結果は、まぁそんなモノかもしれませんね。

最後にエクスポートしてみましょう。

6. エクスポートの実行

エクスポートの手順もほとんど同じです。 Sentencepiece のモデルを差し替えるコードを用意して、

%%bash
cat <<EOF > set_vocab.py
from t5.models import utils
from seqio import vocabularies
from t5.data.utils import DEFAULT_EXTRA_IDS
SPM_PATH = "gs://somewhere/t5/sentencepiece/wikipedia_20190301_ja_v003.model"

def get_vocabulary(mixture_or_task_name=None):
  return vocabularies.SentencePieceVocabulary(SPM_PATH, DEFAULT_EXTRA_IDS)

utils.get_vocabulary = get_vocabulary
EOF

これを --module_import で流し込んでエクスポートすれば OK です¹⁸。

!export PYTHONPATH=${PYTHONPATH}:. && \
  \
  FINE_TUNED_MODEL_DIR='gs://somewhere/t5/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/snow' && \
  \
  echo "FINE_TUNED_MODEL_DIR=$FINE_TUNED_MODEL_DIR" &&\
  \
  t5_mesh_transformer \
    --model_dir="$FINE_TUNED_MODEL_DIR" \
    --checkpoint_mode="specific" \
    --checkpoint_steps="525288" \
    --module_import="set_vocab" \
    --use_model_api \
    --mode="export_predict" \
    --beam_size=4 \
    --gin_param="MtfModel.batch_size=('tokens_per_batch', 40)" \
    --gin_param="MtfModel.sequence_length = {'inputs': 40, 'targets': 40}" \
    --export_dir="./export" 2>&1 | tee export.log

# FINE_TUNED_MODEL_DIR=gs://rddl-nlp/t5/0.8.1/wikipedia_20190301.ja_v003/moe_ex8_hidden3072_cap1.25_drop0.1/snow
# 2021-06-29 07:42:55.800367: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
# ...
# I0629 07:43:45.246580 139821364934528 builder_impl.py:426] SavedModel written to: ./export/temp-1624952581/saved_model.pb

指定した ./export 配下にタイムスタンプでディレクトリが作成され、モデルがエクスポートされます。 エクスポートされたモデルのインタフェースを確認しておきましょう。

!saved_model_cli show --dir ./export/1624952581  --tag_set serve --signature_def serving_default
# The given SavedModel SignatureDef contains the following input(s):
#   inputs['inputs'] tensor_info:
#       dtype: DT_STRING
#       shape: (-1)
#       name: inputs:0
# The given SavedModel SignatureDef contains the following output(s):
#   outputs['inputs'] tensor_info:
#       dtype: DT_STRING
#       shape: (-1)
#       name: strided_slice_1:0
#   outputs['outputs'] tensor_info:
#       dtype: DT_STRING
#       shape: (-1)
#       name: strided_slice_2:0
# Method name is: tensorflow/serving/predict

大丈夫そうですね。あとはこのモデルを Tensorflow Serving にデプロイするだけです。

7. おわりに

今回は Switch Transformer についてご紹介しました。 T5 を最初に触ったときに苦労したかいがあって、 gin で設定をいじるだけで意外と簡単に動作させることができました（のつもりです）。

Google の1兆6千億パラメータのモデルは MoE のアクティベーションに（おそらく） Relu を使ったようですが、実験の中では T5 1.1 で使用した GEGLU も試したようです。 T5 1.1 はオリジナルの T5 よりも性能が良いので、性能を追いかけるなら T5 1.1 準拠の Switch Transformer を作ってみてもよいかもしれませんね¹⁹。

次回はどうしましょうね。じつは先日、 T5 で「文字化けするんですが。。。」と言われ、確認してみたら “UNK” だったということがありました。 有限の語彙数をどこまでレアな文字に割り当てるかというトレードオフの問題で悩ましいなと。そんな訳で次回はトークナイザを使わない ByT5 と Charformer の話にしてみようかなと思ってます。