Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective

要約

一般的なコサイン学習率のスケジュールはステップの総数に依存するため、現在、言語モデルをトレーニングするには、固定のコンピューティング バジェットを事前に決定する必要があります。
対照的に、Warmup-Stable-Decay (WSD) スケジュールでは、一定の学習率を使用して、事前に指定されたコンピューティング バジェットなしで原理的に無期限に継続できる反復のメイン ブランチを生成します。
その後、任意のコンピューティング予算があれば、いつでも適切なタイミングで、急速に減衰する学習率でメイン ブランチから分岐して、強力なモデルを生成できます。
経験的に、WSD は非伝統的な損失曲線を生成します。損失は、安定段階では上昇したままですが、減衰段階では急激に減少します。
この現象を説明するために、事前トレーニング損失は、底に川がある深い谷に似た川の谷の風景を示していると推測します。
この仮定に基づいて、安定段階では、反復は高い学習率により大きな振動を受けますが、それでも反復は川に沿って迅速に進行することを示します。
減衰フェーズでは、学習率が急速に低下するため、反復の振動が最小限に抑えられ、反復が川に近づき、真の最適化の進捗が明らかになります。
したがって、持続的な高い学習率のフェーズと急速に減衰するフェーズは、それぞれ川と山方向の進歩に関与しており、どちらも重要です。
私たちの分析は、経験的観察と一致する現象を予測し、この状況が単純なバイグラム データセットでの事前トレーニングから出現できることを示しています。
この理論に触発されて、以前のチェックポイントの減衰フェーズを再利用し、減衰したチェックポイントから再開するメイン ブランチを 1 つだけ保持する WSD のバリアントである WSD-S を紹介します。
WSD-S は、0.1B から 1.2B までのパラメーターに対して 1 回の実行でさまざまなコンピューティング バジェットにわたって複数の言語モデル チェックポイントを取得する点で、経験的に WSD およびサイクリック コサインよりも優れています。

要約(オリジナル)

Training language models currently requires pre-determining a fixed compute budget because the typical cosine learning rate schedule depends on the total number of steps. In contrast, the Warmup-Stable-Decay (WSD) schedule uses a constant learning rate to produce a main branch of iterates that can in principle continue indefinitely without a pre-specified compute budget. Then, given any compute budget, one can branch out from the main branch at a proper at any time with a rapidly decaying learning rate to produce a strong model. Empirically, WSD generates a non-traditional loss curve: the loss remains elevated during the stable phase but sharply declines during the decay phase. Towards explaining this phenomenon, we conjecture that pretraining loss exhibits a river valley landscape, which resembles a deep valley with a river at its bottom. Under this assumption, we show that during the stable phase, the iterate undergoes large oscillations due to the high learning rate, yet it progresses swiftly along the river. During the decay phase, the rapidly dropping learning rate minimizes the iterate’s oscillations, moving it closer to the river and revealing true optimization progress. Therefore, the sustained high learning rate phase and fast decaying phase are responsible for progress in the river and the mountain directions respectively, and are both critical. Our analysis predicts phenomenons consistent with empirical observations and shows that this landscape can emerge from pretraining on a simple bi-gram dataset. Inspired by the theory, we introduce WSD-S, a variant of WSD that reuses previous checkpoints’ decay phases and keeps only one main branch, where we resume from a decayed checkpoint. WSD-S empirically outperforms WSD and Cyclic-Cosine in obtaining multiple language model checkpoints across various compute budgets in a single run for parameters scaling from 0.1B to 1.2B.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google