パディングは計算を捨て、パッキングは文書の境界をぼかす

一言でいうと

事前学習データは「トークンの長い列」にして、バッチに切って与えます。MiniMindは文書1つを1行に置き、最大長までパディングします。単純で文書が混ざりませんが、文書が短いと計算の大半がパディングに捨てられます。複数の文書をつなげて隙間なく埋めるパッキングは、捨てる計算がない代わりに、1つのウィンドウの中で前の文書が後ろの文書から見えてしまいます。このモジュールでは、私たちのコーパスで2つの方式のコストを数字で測ります。

なぜ必要なのか

モデルは(배치, 길이)の形の整数テンソルを受け取ります(プレースホルダーはバッチ数と長さです)。文書の長さはまちまちなのにテンソルは四角くなければならないので、どこかで長さをそろえる必要があります。方法は2つです。

パディングは、文書1つを1行に置き、足りない分をパディングトークンで埋める方法です。MiniMindのPretrainDatasetが、このようにしています。

tokens = tokenizer(text, max_length=self.max_length - 2, truncation=True).input_ids
tokens = [bos_token_id] + tokens + [eos_token_id]
input_ids = tokens + [pad_token_id] * (self.max_length - len(tokens))
labels = input_ids.clone(); labels[input_ids == pad_token_id] = -100

パディング位置のラベルは-100なので損失には入りませんが、計算は同じように行います。行列の積は、パディングかどうかを知りません。私たちのコーパスの文書は平均33トークンですが、最大長を128にすると、計算の70%以上がパディングに使われます。逆に最大長を短くすると、長い文書が切れてしまいます。MiniMindのREADMEも、このトレードオフを書いています。短いサンプルはパディングで計算を捨て、長いサンプルは切れて情報を失います。そのため、データごとに推奨のmax_seq_lenを別に書いています。

パッキングは、文書ごとに[bos] … [eos]で包んで1本の長い列につなぎ、その列を固定長で切る方法です。捨てる位置はありません。その代わり、1つのウィンドウに文書が3–4個入り、因果マスクは「前のすべてのトークン」を見せるので、後ろの文書のトークンが前の文書を見ます。ほとんどの事前学習は、これを受け入れています。[eos]と[bos]が境界を教えてくれて、モデルが境界の向こうを無視するすべを学ぶからです。境界を越えられないようにするには、文書ごとのアテンションマスクを別に作る必要があり、その分だけ実装が重くなります。

どう動くのか

パッキングした結果は、整数の配列1つです。語彙が65,536より小さければuint16で足りて、int64の4分の1の容量で済みます。学習ループは、この配列から任意の位置を選び、その位置から길이の分だけ切り出してバッチを作ります(プレースホルダーは長さです)。

ix = torch.randint(0, len(train) - seq, (batch,), generator=g)
x = torch.stack([train[i:i + seq] for i in ix])
loss = model(x, labels=x).loss      # 한 칸 미는 일은 모델 안에서

MiniMindモデルは、labelsとして入力と同じテンソルを受け取り、内部でlogits[..., :-1]とlabels[..., 1:]を合わせて、次トークン予測の損失を出します。そのため、データ側で入力と正解を別々に1つずらしておく必要はありません。

検証データは文書単位で切り分ける必要があります。1つの配列を前の95%・後ろの5%で切ると、1つの文書が2つに分かれて両側にまたがることがあり、同じ文書が学習と検証の両方にあると、検証損失は暗記したものを測ります。切り分けたあとも、同じ文章が学習側にそのまま残っていないか、一度確認します。

現場での姿

短い質疑応答やチャットの記録のように文書が短いデータをパディングで学習すると、GPUは忙しく動いているのに、損失はゆっくりしか下がりません。使用率は高くても、実際に学習するトークンが少ないからです。このとき、1学習ステップで処理した本当のトークン数を数えてみると、原因がすぐに見えます。逆に、パッキングに変えたあと、モデルが文書の境界を越えて見当違いの内容を続けて書きはじめたなら、まず[eos]を正しく付けたかを確認します。

学習予算もトークンで数えます。「学習ステップ数」は、バッチサイズと長さによって意味が変わるので、1学習ステップが処理するトークン数と、コーパス全体のトークン数から、何学習ステップで1周(エポック)になるかを計算しておかないと、損失曲線の凹凸を読めません。

MiniMindのオリジナルとこのコースの違い

MiniMindは、HuggingFaceのdatasetsでjsonlを読み、サンプルごとに__getitem__でその都度トークナイザーを動かします。1.2GBのコーパスを前もってトークンに変えておかなくてもよい代わりに、学習している間、CPUがずっとトークンを切り続けます(そのため、デフォルトのnum_workersが8です)。このコースはコーパスが2MBにも満たないので、一度にすべて切ってuint16の配列1つとして保存し、学習ループは、その配列からウィンドウを切り出して使うだけです。大規模な事前学習でも、この方式(前もってトークン化してバイナリファイルにする)がよく使われます。データを何周も見る間、同じトークン化を繰り返さずに済み、ファイルをメモリマップして必要な部分だけ読めるからです。

次のラボですること

基準のトークナイザーで文書の長さを測り、MiniMind式のパディングがどれだけ捨てるかを数えます。学習・検証コーパスを文書単位でパッキングしてuint16の配列として保存し、検証文書が学習側に漏れていないか、パッキングしたウィンドウでトークンの何%が前の文書を見られるか、1周が何学習ステップかを計算します。