Token Weighting: トークン重み付け

概要

Token Weighting（トークン重み付け）は、訓練データに含まれる様々な長さの出力を持つタスクにおいて、損失関数の寄与度を調整する手法である。Molmo2 の訓練では、単一トークンの選択肢問題から 4,000 トークン以上の長いビデオキャプションまで、多様な出力長を持つデータが混在している。

問題設定

長い出力による損失の支配

訓練データに出力長の大きな偏りがある場合、以下の問題が発生する:

トークン数の偏り: 長い出力を持つタスク（例: ビデオキャプション）は、サンプリング頻度が低くても、訓練時の損失トークンの大部分を占めてしまう
短い出力タスクの劣化: 多肢選択問題や短い回答を要求するタスクの性能が低下する
タスクバランスの崩壊: モデルが長い出力を生成するタスクに過度に最適化され、他のタスクでの性能が犠牲になる

例: 極端なケース

1 サンプルのビデオキャプション（4,000 トークン）と 100 サンプルの多肢選択問題（各 1 トークン）を同じバッチで訓練する場合、ビデオキャプションが損失全体の約 97.5% を占めることになる（4,000 / (4,000 + 100) ≈ 0.975）。

Molmo2 の重み付け戦略

Molmo2 では、タスクの種類と出力長に応じて、各サンプルの損失に対する重みを調整する。

重み付けスキーム

タスク種別	重み	理由
ビデオキャプション	0.1	非常に長く密な出力（4,000+ トークン）を生成するため
ポインティング	0.2	座標列挙により長く密な出力を生成するため
その他のタスク	\(\frac{4}{\sqrt{n}}\)	出力長 \(n\) に応じた適応的な重み付け

数式による定義

サンプル \(i\) の損失重み \(w_i\) は以下のように定義される:

\[ w_i = \begin{cases} 0.1 & \text{if task is video captioning} \\ 0.2 & \text{if task is pointing} \\ \frac{4}{\sqrt{n_i}} & \text{otherwise} \end{cases} \]

ここで、\(n_i\) はサンプル \(i\) の回答トークン数である。

適応的重み付けの効果

その他のタスクに対する \(\frac{4}{\sqrt{n}}\) という重み付けは、以下の特性を持つ:

短い出力: \(n = 1\) のとき \(w = 4.0\)、\(n = 16\) のとき \(w = 1.0\)
中程度の出力: \(n = 100\) のとき \(w = 0.4\)
長い出力: \(n = 400\) のとき \(w = 0.2\)

この平方根による減衰により、長い出力を持つサンプルの影響を抑えつつ、完全に無視することを避ける。

重み付けの直感

平方根による重み付け \(\frac{4}{\sqrt{n}}\) は、出力長が 4 倍になると重みが半分になる特性を持つ。これにより、長い出力と短い出力の間でバランスの取れた学習が可能になる。

例えば:

1 トークン出力: 重み 4.0
4 トークン出力: 重み 2.0
16 トークン出力: 重み 1.0
64 トークン出力: 重み 0.5

効果

長短出力のバランス改善

Token Weighting の導入により、以下の効果が得られる:

多様なタスクでの性能維持: 短い回答を要求するタスク（多肢選択問題など）と長い出力を生成するタスク（ビデオキャプションなど）の両方で良好な性能を達成する
訓練の安定性向上: 損失が特定のタスクに支配されることを防ぎ、より安定した訓練が可能になる
効率的なデータ活用: 出力長が異なる多様なタスクを単一のモデルで効果的に学習できる

実装上の考慮事項

Token Weighting は損失計算時に適用される:

\[ \mathcal{L} = \frac{1}{B} \sum_{i=1}^{B} w_i \cdot \mathcal{L}_i \]

ここで:

\(B\) はバッチサイズ
\(\mathcal{L}_i\) はサンプル \(i\) の未調整損失
\(w_i\) はサンプル \(i\) の重み

この重み付けは、各サンプルの損失を個別に計算できるため、実装が容易であり、既存の訓練パイプラインに容易に統合できる。