BERTをfine-tuningする際のTips [自分用メモ]

AI SHIFT様が作成されたブログ記事が分かりやすかったので、筆者としての解釈を自分用メモとして残すことにする。

Kaggleで学んだBERTをfine-tuningする際のTips①〜学習効率化編〜 | 株式会社AI Shift (ai-shift.co.jp)

TIPS1: 混合精度(Mixed Precision)の利用

Tensorコアが導入されている最新のGPU(NVIDIA A100など)を用いることで学習時間を短縮

論文は百度とNVIDIAが発表したMixed Precision Training

要するに、深層学習に最適なGPUを選択しましょうね、ということ

TIPS2: 文章の切り詰め

当たり前のことだが、それぞれの文に含まれている単語数はばらばらなので、BERTに入力する際には[PAD]トークンを用いて下図のようにパディングする必要がある。

Fixed Padding
Speeding up Transformer w/ Optimization Strategies | Kaggleより引用

このとき、文の長さが大体同じデータセットであれば、気にする必要はないが、長い文と短い文が同時に含まれているデータセットの場合、穴埋めされる量に偏りが生じてしまう。より具体的に言うと、短い文の場合、[PAD]トークン(穴埋め)が増えて、結果、情報量が少ない割に、GPUメモリを圧迫するサンプルが生じてしまうことになる。

そこで、文の長さをそろえることによって、文章を切り詰めるUniform Length Batchingという手法が提案されている。上の図と比べると、[PAD]される領域が減少していることがみてとれ、学習時間の削減が期待される。しかし、バッチごとに含まれる文の数にばらつきが生じてしまう可能性があるため、実際にこの手法を用いる際は注意されたし。

Uniform Length Batching
Speeding up Transformer w/ Optimization Strategies | Kaggle

TIPS3: 勾配累積(Gradient Accumulation)

BERTなどのパラメータ数が多いモデルを扱う際は、往々にしてメモリ消費量が大きく、実際にfine-tuningすると往々にしてOOMエラーが発生する。このとき、多くの開発者はバッチサイズを小さくする調整を行うであろう。しかし、小さなバッチサイズだと、学習に余計な時間がかかったり、学習が不安定になることも考えられる。そこで、勾配累積:小さいバッチで計算した重みを、複数回分ためてから平均をとり、それを用いてモデルのパラメータを更新する、という手法を取り入れる。

AI SHIFT様の記事より簡易的なコードを引用させていただき、さらに、筆者のコメントも加えて貼らせていただく。

ITERS_TO_ACCUMULATE = 10 # 累積数

for epoch in epochs:
    for i, (input, target) in enumerate(data):

        output = model(input)
        loss = loss_fn(output, target)
        loss = loss / ITERS_TO_ACCUMULATE # バッチごとに計算されたロスの平均を蓄積
        loss.backward() # 目的の関数(ここではloss)に対して微分を行った時の勾配を計算。累積数が10なので、loss.grad(勾配)は10回分追加される。
        
        # 累積数(ここでは10回)ごとに1回下記内容を実行
        if (i + 1) % ITERS_TO_ACCUMULATE == 0:
            optimizer.step() # これは、適当な学習率をかけてパラメータを更新するメソッド
            optimizer.zero_grad() # これは、一旦計算した勾配の初期化

コードをみると、実装自体は簡単であることがわかるだろう。通常、optimizer.step()とoptimizer.zero_grad()はバッチごとに1回適用されるのを、累積数(ここでは10回)ごとに1回適用するという形に変更するという実装内容になっている。

BERT,PyTorch,Transformer

Posted by vastee