- RAdamScheduleFree という新しいoptimizerを作りました。
- warmup、learning rate scheduler、訓練終端時刻の指定、全て不要です。
- 安定かつ高速に収束します。多くの場合でAdamやAdamWより強いです。
この Version 1.4 から RAdamScheduleFree が搭載されました。
みんな使ってね。
はじめに
お久しぶりです。はまなすです。
某記事を執筆してからそろそろ2年が経ちそうなことに戦慄を覚えつつ、いまでも新たに反応をいただくこともあり、じんわり嬉しさを覚える年の瀬です。
さて、実務や趣味で機械学習開発に携わっている方々は、皆等しくAdam（またはAdamW）とお友達だと思います。Adam台頭以降、RAdamやAdaBound、AdaBeliefなど、直接的な後継といって差し支えないものだけでも、新たなoptimizerは数多く登場してきました。Lionなどの別系統も[1]。しかし、デファクト・スタンダードと呼べるoptimizerはいまだにAdamです。なぜでしょうか。
ひとつには、Adamにしておけばひとまずは問題ないという安心感がありそうです。もはやある種の枯れた技術となったAdamは、どこの馬の骨ともわからないoptimizerを試すより遥かに信頼できる実績が積み重なっており、とりあえずそれを選んでおくことで「もしかしてこの謎挙動、Adamを使わなかったせいかな…？」みたいな無益な精神の摩耗から我々を保護してくれます。
もうひとつには、Cosine annealing scheduler などの強力なschedulerが同様にデファクト・スタンダードに君臨していることが挙げられるでしょう。Adamと Cosine annealing scheduler の組み合わせは暴力的に汎用性が高く、そもそもoptimizerを変えようという発想を（optimizerジャンキーを除く一般的な）開発者から奪い去っています。それは同時に、本来はschedulerの導入にあたって発生する新たな自由度を飼い慣らす必要性さえをも我々の意識から遠ざけ、「とりあえずこれでいいや」という盲目的な気持ちの増長に一役買っています。本当は、適切な調整によりもっといいモデルが訓練できる余地があるかもしれないのに、です。
でも、タスク設計やデータ整備、モデルアーキテクチャなどのより本質的な部分に思考リソースを割きたい開発者にとって、optimizerやschedulerの調整にまで気を配るのは正直面倒なはずです。
では、そもそもそんな学習率スケジューリングなんて全く要らないoptimizerがあるとしたら？
ScheduleFree という新パラダイム
2024年春、ScheduleFree という新たなoptimizerシリーズがMetaより突如提案されました。
詳細は理論背景の章に譲るとして、これこそがまさに「学習率スケジューリングなんて全く要らないoptimizer」の正体そのものです。初期リリースではSGDとAdamWの ScheduleFree version が実装されており、論文内の実験でもこれら2種類における性能の高さが実証されています。
簡潔にまとめると、ScheduleFree、特にその主提案である AdamWScheduleFree によって、次のようなことが実現されました。
- ［手軽］ いかなるlearning rate schedulerも不要。訓練終端時刻の指定も不要。
- ［性能］ scheduler込みで調整したAdamWより安定で高速。収束品質も同等かそれ以上。
The Road Less Scheduled: Figure 5｜既存のoptimizerに learning rate schedule を組み合わせたものと ScheduleFree の比較。上4つがSGD、下4つがAdamWによる実験結果。画像分類や画像の継続事前学習、言語翻訳や言語モデルなど、種々のタスクやモデルアーキテクチャにて強力な優位性と安定性を実証。
The Road Less Scheduled: Figure 7｜AlgoPerf challengeに基づく、音声を含む様々なモダリティやタスクにおける検証結果。ベースラインであるNesterov AdamWと比較して、同様に優位性と安定性を実証。
このように旧来のベストプラクティスを次々と駆逐した AdamWScheduleFree でしたが、その影で「warmup stepsを指定する必要性」だけが若干の手間として残されました。
学習率のwarmupはいわずとしれた重要な調整項目で、とりわけAdamWのように勾配の二次指数移動平均を用いるoptimizerは、訓練初期の安定性に無視できない影響を受けます。これも「とりあえず設定しておけばいいハイパラ」の筆頭で、実際のところ、デフォルト値以外をきちんと探索する気概や余裕のある開発者は限られるのではないかという実感があります。
そこで登場するのが古株のRAdamです。RAdamは、AdamWでは職人芸的に調整するしかない最適なwarmup stepsと同等性能を、ハイパラ設定なしに実現してくれます。素晴らしいoptimizerです。あとはお分かりですね。
RAdamの ScheduleFree version があれば、我々はすべての苦行から解放されます。
なので作りました。
これ使っときゃOK、を次の水準へ
なにはともあれインストール。とっても簡単ですね。
pip install schedulefree
なんらかのパッケージマネージャをお使いの場合でも、普段通りで大丈夫です。
rye add schedulefree
rye sync
それでは ScheduleFree の使い方を見ていきましょう。といっても、基本的には従来のoptimizerとなんら変わりません。差分は以下の2点です。
- scheduler が不要になること
-
optimizer.eval()とoptimizer.train()を適切なタイミングで呼ぶこと
例えば、AdamWとなんらかのschedulerを使う場合のよくあるモックを考え、そこからの差分として骨子を表現してみると、以下のような感じになるかと思います[2]。
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
+from schedulefree import RAdamScheduleFree
from torch import Tensor
-from torch.optim.lr_scheduler import CosineAnnealingLR, LRScheduler
from torch.utils.data import DataLoader
class MyModel(nn.Module):
r"""
Awesome neural network.
"""
def __init__(self, *args, **kwargs):
...
def forward(self, x: Tensor) -> Tensor:
...
return y
class MyDataset(nn.Module):
r"""
Awesome dataset.
"""
def __init__(self, *args, **kwargs):
...
def __getitem__(self, index: int) -> tuple[Tensor, Tensor]:
...
return data, target
def train(
model: nn.Module,
trainloader: DataLoader,
optimizer: torch.optim.Optimizer,
- scheduler: LRScheduler,
device: torch.device,
):
r"""
Train the model for one epoch.
"""
model.train()
+ optimizer.train()
for batch in trainloader:
data, target = map(lambda x: x.to(device), batch)
optimizer.zero_grad()
pred = model(data)
loss = F.mse_loss(pred, target)
loss.backward()
optimizer.step()
- scheduler.step()
def validation(
model: nn.Module,
validloader: DataLoader,
+ optimizer: torch.optim.Optimizer,
device: torch.device,
):
r"""
Evaluate model performance on the validation set.
"""
model.eval()
+ optimizer.eval()
with torch.inference_mode():
for batch in validloader:
data, target = map(lambda x: x.to(device), batch)
pred = model(data)
# update metrics
...
def save(
save_dir: Path,
model: nn.Module,
optimizer: torch.optim.Optimizer,
- scheduler: LRScheduler,
):
r"""
Save model and training states to a checkpoint file.
"""
torch.save(
{
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
- "scheduler": scheduler.state_dict(),
},
save_dir / "checkpoint.pt"
)
def load(
load_dir: Path,
device: torch.device,
model: nn.Module,
optimizer: torch.optim.Optimizer,
- scheduler: LRScheduler,
):
r"""
Load model and training states from a checkpoint file.
"""
pkg = torch.load(load_dir / "checkpoint.pt", map_location=device)
model.load_state_dict(pkg["model"])
optimizer.load_state_dict(pkg["optimizer"])
- scheduler.load_state_dict(pkg["scheduler"])
def main():
# =========================== #
# Training settings #
# =========================== #
parser = argparse.ArgumentParser()
parser.add_argument('--epochs', type=int, default=1)
parser.add_argument('--save', type=str, default=None)
parser.add_argument('--load', type=str, default=None)
...
args = parser.parse_args()
device = ...
# ====================== #
# Preparations #
# ====================== #
model = MyModel(...).to(device)
params = model.parameters()
- optimizer = torch.optim.AdamW(params, lr=1e-4, betas=(0.9, 0.999))
+ optimizer = RAdamScheduleFree(params, lr=1e-4, betas=(0.9, 0.999))
- scheduler = CosineAnnealingLR(optimizer, ...)
if args.load is not None:
load(
Path(args.load),
device,
model,
optimizer,
- scheduler,
)
trainloader = DataLoader(MyDataset(...), ...)
validloader = DataLoader(MyDataset(...), ...)
# ================== #
# Training #
# ================== #
for epoch in range(args.epochs):
train(
model,
trainloader,
optimizer,
- scheduler,
device
)
validation(
model,
validloader,
+ optimizer,
device
)
if args.save is not None:
save(
Path(args.save),
model,
optimizer,
- scheduler,
)
if __name__ == '__main__':
main()
はい、schedulerがことごとく抹消されていることと[3]、optimizerに対してtrainやevalというメソッドが呼ばれていることがわかりますね[4]。
このoptimizerの状態切り替えが何をやっているかは、「ScheduleFree が何をしているか」と「実際の実装でそれがどう実現されているか」の両方を理解しないと正確には理解できないので、一旦は『modelとセットでoptimizerの状態も切り替える必要がある』とだけ覚えておいてください。
また、上例ではvalidationのすぐ後にsaveがあるので気にしなくてよかったのですが、『modelやoptimizerの保存・読み込み時は、optimizerが eval mode になっている必要がある』ことも頭の片隅に置いておいてください[5]。例えば、checkpointの保存が訓練ループの中にあるような場合には、torch.saveの前にoptimizer.eval()を呼んでくださいね[6]。
さて、RAdamScheduleFree を実践で使用するための話は、ここまでで全て終わりました。他の多くのoptimizerと同様、中でなにをしているか知らなくても使えます。重要なのは、今回の提案により『warmup、learning rate scheduler、訓練終端時刻の指定、全てが不要になる』ということです[7]。これにより、我々はもっと本質的な開発や調整に時間を割くことができるようになります。
ところで、詰まるところ今回の取り組みスコープは既存の強いoptimizerの組み合わせでしかなく[8]、特にこれで論文を書いたりするわけでもないので、新規に網羅的な性能実験などはおこなっていません。つまり、皆さんにとっては依然として「どこの馬の骨ともわからないoptimizer」の類ではあるわけですが、それをあなたにとっての新しい「これ使っときゃOK」にするかどうかは、あなたの好奇心次第です。この記事の目的は、冒頭に述べたもうひとつのほうの懸念を潰すことで、その背中を押すことにもありました[9]。
少なくとも、私の使用感では RAdamScheduleFree はかなり強いです。なにより超楽。
皆さんもぜひ使ってみてくださいね。こんな問題でもうまくいったよ、こんな問題だと期待薄だったよ、みたいな感想も歓迎しています。気が向いたらお気軽にお寄せください。
──さて、沼にはまる覚悟はありますか？
理論背景
ここからは関連する理論背景について話していきます。怪しげな ScheduleFree の概要を理解したい方には多分面白い内容になると思いますが、もともと本記事の想定読者に据えている「optimizerの選定や調整なんて面倒すぎるよ〜」という方々にとっては特に有益な情報ではないと思うので、そのあたりの温度感をご理解いただいた上で読み進めていただければ幸いです。
以降、かなり玄人向けです。主に次の三本立てでお送りします。
- ScheduleFree のお気持ちを理解する
- RAdamを思い出す
- そして RAdamScheduleFree へ
ScheduleFree のお気持ちを理解する
いきなり二次指数移動平均を用いる最適化手法を考えると話がややこしくなるので、モメンタムも考えない最もシンプルなSGDを出発点に ScheduleFree の挙動を確認していきましょう。
単純なSGDと、パラメータの反復平均化
モメンタムのないSGDは、更新されるパラメータ系列
さて、このような愚直な勾配降下は準最適な（つまり最適とはいえない）結果に落ち着くことが広く知られています。そこで、このようにして更新されるパラメータ系列
各時刻で勾配を計算するための（すなわち訓練時に評価される）パラメータは
計算してみるとわかりますが、
ところで、Polyak averaging が想定する関数
話は変わり、比較的最近提案された平均化手法として、Primal averaging というものも存在します[10]。Primal averaging は見かけ上 PR averaging とそっくりですが、勾配を評価するパラメータが
実は、うまく定数変換をすることで、Primal averaging は注意深く学習率スケジューリングされたモメンタム付きSGDと同一視できることが示されます。その意味で、機械学習最適化における超基礎たるモメンタム付きSGDは、勾配の指数移動平均ではなく、純粋なSGDにおけるパラメータの反復平均化という世界へ接続するのです。
これを踏まえて、改めてそれぞれの平均化の表式を眺めてみましょう。PR averaging では
パラメータ更新は速いが実践に弱い PR averaging と、収束の遅さは拭えないが「注意深く学習率スケジューリングされたモメンタム付きSGD」の形で有効性が実証されている Primal averaging という図式が揃いました。綺麗なトレードオフです。
こういうとき、両者のいいとこどりをしたいのが人情というものです。
パラメータ反復平均化の補間としての SGDScheduleFree
PR averaging と Primal averaging を内挿するような定数
安定性と更新速度のいいとこどりを実現する実践的な値として、論文では
ScheduleFree の収束性［最も抽象的で、難しい］
さて、このようにして導出された SGDScheduleFree は、学習率スケジューリングをしなくても優秀な収束性を持つことが示されています。論文内で最初に示される性質は、Lipschitz連続な非平滑凸関数設定を考えたとき、どのような
Lipschitz連続な凸関数は、機械学習の確率的最適化手法を考える際によく出てくる問題のクラスです。先ほども出てきましたね。平たくいうと『どの2点を取っても、その間の関数の変化の大きさが距離に比例して抑えられている凸関数』のことです。急峻な変化を持つ箇所がどこにも存在しない、基本的にはなだらかなお椀型の形状を持ち、理論的な解析がしやすい関数のクラスと考えればよいでしょう。例えば、どのような
また最悪ケース最適とは、理論的に可能な最高の最悪ケース収束率を達成していることを意味します。最悪ケース収束率とは『それ以上遅い収束はあり得ない』ことを表現する収束性指標で、例えば最適化の終端時刻
いま、
論文では特に記載がありませんが、次のように変数を消去するともっとわかりやすいかもしれません。
ところで、Lipschitz連続な平滑凸関数の最悪ケース最適速度は
結局 SGDScheduleFree がどういう性質を持つかというと、「Lipschitz連続な非平滑凸関数において、どんな
論文ではこの議論をさらに拡張し、
後学のため、定式化も確認しておきましょう。独立同分布な変数系列
このとき、凸関数
この時点で
そういうわけで、
色々と小難しい言葉を並べましたが、これまではschedulerを用いて時刻ごとに変化する学習率
直感的な理解をすると、
先に説明したように推論時は
このイメージを具象化してみましょう。論文の行間を埋める作業として、SGDScheduleFree において
興味深いことに、綺麗に線形減衰する寄与度を導くことができましたね。Primal averaging や SGDScheduleFree に対しても同様の議論をおこなうと下図が得られます。
The Road Less Scheduled Figure 4｜ある時点のスナップショット（総訓練時間
ここから分かるのは、SGDScheduleFree が内部的に Linear decay scheduler に類する挙動を実現しているということです。学習率スケジューリングを陽に設定せずとも、自動的に。このような観点からも、ScheduleFree の ScheduleFree たる所以が垣間見えます。
余談ですが、このように優秀な収束性を持つ ScheduleFree は、実践上も非常に面白い挙動を示します。下図に示すように、既存のoptimizerとschedulerを組み合わせ、異なるハイパラチューニングをして複数実験した結果を集めると、その訓練の行先が ScheduleFree の訓練曲線にほぼ一致するのです。これを論文では「訓練時間と評価損失のパレート・フロンティア」と表現しています。
The Road Less Scheduled Figure 1｜ScheduleFree（黒線）が、1回の試行で「訓練時間と評価損失のパレート・フロンティア」をおおよそなぞる挙動。左がSGD、右がAdamW。赤線はそれぞれのoptimzierにcosine schedulerを組み合わせ、ハイパーパラメータを変えて試行した複数の実験を表す。ScheduleFreeは、既存のoptimizerとschedulerの組み合わせが実現し得る「訓練時間と評価損失のトレードオフ」の曲線をなぞるようにしながら、任意の時間訓練し続けられるという特性が読み取れる。
多変量同時最適化において各変数の最適性バランスを取った解はパレート解と呼ばれ、一意に定まらないパレート解の集合が描く曲線（あるは一般に超曲面）はパレート・フロンティアと呼ばれますが、ここでは「より良い評価損失を得るためには十分な訓練終端時刻（とそれに依存して学習率を減少させるscheduler）が必要」という既存システムのトレードオフを指してこのように表現しているようです。図から分かるように、ScheduleFree はこのパレート・フロンティアを1回の試行で悠々と撫でていくかのようです。トレードオフにより実現されるあらゆるパレート解を辿りながら、好きなだけ長い時間訓練し続けられるわけですね。ここに、ScheduleFree が訓練終端時刻
メモリ効率の良い SGDScheduleFree の実装
話を SGDScheduleFree に戻しましょう。これを愚直に実装すると、
ScheduleFree のリポジトリを参照すると、SGDScheduleFreeReference というクラスにこれまでの定式化を素直に反映した実装が組み込まれており、理論と実装の対応を丁寧に追うことができます。これはあくまで参考用実装で、使用すると先に述べたように余計なメモリ消費を招きます。一方、通常利用が想定される SGDScheduleFree クラスは、実装を効率化することでメモリ消費量をモデルパラメータの2倍、つまり通常のモメンタム付きSGDと同じレベルまで削減しています。その実装が何をしているかを式に起こすと、次のようになります。
つまり、
実践の説明で「optimizer.eval()を呼んでくださいね」と話した意味が、ここに繋がります。つまり、optimizerが train mode のときはモデルパラメータとして optimizer.train()は実装されています。
ちなみに、
AdamWScheduleFree への拡張
SGDScheduleFree を AdamWScheduleFree に拡張するのはそんなに難しい話ではなく、AdamWにて提案された以下の2点、
- 二次指数移動平均による適応的ステップサイズ
- weight decay
\lambda
および、warmup steps
このとき、適応的ステップサイズ
さて、このアルゴリズムは、ScheduleFree のリポジトリにて AdamWScheduleFreePaper として実装されています。わざわざPaperという接尾辞がついているのはなにやら不穏ですね。これは、著者らがその後、
以上で、AdamWScheduleFree が何をしているかを全て追いかけることができました。同時に、記事の導入で触れたように warmup steps
RAdamを思い出す
RAdamは5年以上前に提案された手法なので、既に有用な日本語解説記事がいくつもあります。
RAdamが何をしているかを簡潔に述べると、「二次指数移動平均係数
式を見ればわかるように、
この
この定式化であれば、
そして RAdamScheduleFree へ
ここまでで、AdamWScheduleFree の warmup steps
前節でRAdamの実装にはいくつかの派閥が存在するという話をしましたが、今回の私の実装では
先述のように、最初の数ステップでSGD的な挙動をするか否かは全体の最適化においては瑣末な差異でしかなく、立ち上がりの安定性を担保することこそが最優先です。この判断に基づき、SGDフェーズを撤廃するような実装方式を取りました。この挙動は RAdamScheduleFree の初期化引数 silent_sgd_phase で制御しています。デフォルト値Trueのままにしておくのを推奨します。
以上、RAdamScheduleFree にまつわる理論背景説明、全て終了です。お疲れ様でした。
結び
本記事では、Adamに代わって新たに普段遣いしてほしいoptimizerとして RAdamScheduleFree を提案しました。概要に始まり、置き換えると何が嬉しいのかを実践例ベースで解説するとともに、内部挙動に興味のある方向けに要素技術の背景についても説明を施しました。
RAdamScheduleFree は汎用的で、強く、手軽です。本記事が、皆さんが取り組まれている様々な機械学習プロジェクトの開発効率をちょっとでも高めるエッセンスになれることを願い、筆を置きます。それではまた、いずれどこかで。
参考文献
- The Road Less Scheduled [Defazio et al., NeurIPS, 2024]
- On the Variance of the Adaptive Learning Rate and Beyond [Liu et al., ICLR, 2020]
- Momentum via Primal Averaging: Theoretical Insights and Learning Rate Schedules for Non-Convex Optimization [Defazio, 2020]
- Optimal Rates in Convex Optimization [Tibshirani et al., 2014]
-
optimizerの提案史は適当に眺めるだけでも群雄割拠で、いまや古典的なものだけでも SGD、モメンタム付きSGD、NAG、AdaGrad、RMSProp、AdaDelta など数多く、Adam以降はもはや全容を把握するのは不可能なほど溢れています。Adamと地続きなものだけでも AdaMax、AdamW、NAdam、AMSGrad、RAdam、AdaBound、AMSBound、AdaBelief、最近だとADOPTや Cautiousシリーズ などが挙げられますし、視野を広げれば Santa、Eve、YellowFin、Lion、最近だと強力なものとして Shampoo、Distributed Shampoo、SOAP などが提案されています。 ↩︎
-
余談ですが、コードサンプルって MyModel や MyDataset などのトイクラスがどこからともなく現れるものが多いですよね。コンパクトさを重視するなら絶対そのほうがいいのですが、私は「その MyModel どこからimportしたんだよ」みたいな違和感で萎える性分なので、気持ち悪くない範囲であえて丁寧めに書いています。深層学習のコードってどうしても準備やお作法が結構ボリューミーですが、その辺を端折ると実践に当てはめる際のヒントとしては不十分になって迷わせてしまうこともあるので……。 ↩︎
-
余談ですが、これまではschedulerが直接
optimizer.param_groups[i]["lr"]を書き換えて学習率を操作することが多かったものの、ScheduleFreeでは"lr"属性は変化せず、optimizer.param_groups[i]["scheduled_lr"]に読み取り用の値が格納される様式になっています。TensorboardやW&B等で学習率をモニタリングする際の属性は"scheduled_lr"を使ってくださいね。 ↩︎ -
より詳しい話をすると、ScheduleFree シリーズにはそれぞれ ScheduleFreeClosure という亜種が実装されており、私も RAdamScheduleFreeClosure も含めて実装をしました。こちらはPyTorchのclosureを用いてoptimizerを動かす場合のクラスで、
optimizer.step(closure)を呼ぶとこのmode切り替え相当の処理が内部で自動実行されるので、明示的にユーザーがtrainやevalを呼ぶ必要はなくなっています。ただ、closureを使わない実装の方が一般的であることや、ScheduleFreeClosure は ScheduleFree と比べごく僅かに訓練効率が劣ること、本文に ScheduleFreeClosure についての言及を含めると混乱を招く可能性があることから、脚注にて補足しました。関連の記述は本家のHow to Useをご覧ください。 ↩︎ -
これは、optimizerのmode切り替えによりモデルパラメータ自体が書き換えられるためです。 ↩︎
-
仕様の補足ですが、ScheduleFree 系のoptimizerは初期化時には eval mode になっています。eval mode の時に
optimizer.eval()を呼んだり、逆に train mode の時にoptimizer.train()を呼んでも何も起きないので、重ねがけによる副作用を心配する必要はありません。 ↩︎ -
どこに補足するか迷ったのでこちらに。RAdamScheduleFree はAdamWと同様に decoupled weight decay を採用しているので、必要があれば weight decay を適用しても齟齬なく動作します。命名には陽に含めていませんが、実態としては RAdamWScheduleFree と思っていただいても問題ありません。 ↩︎
-
組み合わせるにあたってちょっとした工夫はあり、それについては本記事の終盤で紹介しています。ただ、RAdamと ScheduleFree という巨人の肩に乗った上でささやかに口笛を吹いたくらいのものなので、「組み合わせるというものでしかなく」という表現でも特に差し支えはありません。 ↩︎
-
背中を押してなんの得があるのだと思われるかもなぁと書いていて思いましたが、自分が使っていて効果が実感できたものを純粋に人にもお勧めしたいときってありますよね。そんな感じです。あとは、業界基準が更新されることが巡り巡って将来の自分のためにもなるかなという淡い感覚と、長らく塗り替えられなかった『optimizerのデファクト・スタンダードに関する全体意識』を刷新する一助になったら超面白いな、という気持ちもあるかもしれません。実際、Adamからの乗り換えを阻む主要因は、個々のoptimizerの性能云々よりもむしろ、乗り換えの面倒さだと思っています。その面倒さを極力排除した上で、既存の面倒さもついでに消し去りますよーみたいなお膳立てができればこの記事としてはいいのかなと考えて執筆しました。 ↩︎
-
PR averaging が1990年前後に提案された手法であるのに対し、Primal averaging は2015年~2020年頃に整理されたものなので、近年の深層学習ブームに照らして再考されたパラメータ平均化手法と言えるかもしれません。 ↩︎
-
実際には無限に長く訓練することはできませんが、浮動小数点の表現精度の限界で実質
x_t
Discussion