スケーリング則

LLMの事前学習では、モデルサイズ N、学習トークン数 D、計算資源 C を大きくすると、検証損失がべき乗的に下がる傾向があります。
このノートでは、実験データからべき乗則を推定し、同一計算資源（isoflops）での最適配分を計算する流れを確認します。

計算資源が固定なら、何を増やすべきかは自明ではない

モデルを大きくするべきか、データを増やすべきか。事前学習ではこの配分の誤りがそのまま高価な遠回りになります。スケーリング則は、その判断を勘ではなく観測値から支えるための道具です。

このノートでは、まず損失がべき乗的に下がるという見方を作り、次に $L_inf$ を含むフィットを試し、最後に同一計算資源での N と D の配分を考えます。目的は式を覚えることではなく、訓練計画の読み方を身につけることです。

読みどころは、損失曲線そのものより配分判断です

きれいな log-log 直線を眺めることが本題ではありません。大事なのは、その直線を使って「この予算ならモデルをこれ以上大きくする意味があるか」を考えられるようになることです。

ここでは合成データを使うが、見たい構造は本物と同じ

実験としては toy ですが、N, D, C の関係を読む練習としては十分です。どこまでが教育用の簡略化で、どこからが実運用でも有効な感覚かを切り分けながら読み進めます。

まずは「一直線に見えるか」より、何を仮定しているかを意識する

べき乗則は万能法則ではなく、ある領域での経験的な近似です。データ品質、アーキテクチャ変更、最適化条件が変われば崩れます。その前提を置いた上で、式を判断材料として使います。

後半では、式をそのまま予算配分の会話へつなぐ

同じ FLOPs でも、巨大モデルに少量データを食べさせるのか、中規模モデルに大量データを食べさせるのかで結果は変わります。ここが実務上いちばん効く読みどころです。

この notebook は発見の再現ではなく、判断軸の習得に寄せている

ここでのフィットは教育用です。ただ、何を増やしたときに何が頭打ちになるか、という視点はそのまま本番の実験計画に持ち込めます。

まずは損失の形を数式で置く

最初に、スケーリング則でよく使われる近似式を置きます。ここでの目的は厳密な理論証明ではなく、「何が頭打ちを表しているのか」を式の中で見分けられるようになることです。

import math
import numpy as np
import matplotlib.pyplot as plt

よく使う近似は次の形です。

L(N, D) ≈ L_inf + a * N^{-alpha} + b * D^{-beta}

$L_inf$ : これ以上は下げにくい下限（不可約損失）
alpha, beta: モデル拡大・データ拡大の効き方

まずは N と D をそれぞれ変えた観測データを作り、alpha, beta を推定します。

# 観測例（教育用の合成データ）
N_million = np.array([30, 60, 120, 240, 480, 960], dtype=np.float64)   # million params
D_billion = np.array([5, 10, 20, 40, 80, 160], dtype=np.float64)        # billion tokens

# 実際の単位へ変換
N_params = N_million * 1e6
D_tokens = D_billion * 1e9

# 生成側の真値（未知だと思って推定する）
L_inf_true = 1.60
A_true, alpha_true = 3.2, 0.37
B_true, beta_true = 2.1, 0.29

# 損失生成では見やすさのため M/B 単位でべき乗を作る
rng = np.random.default_rng(7)
L_of_N = L_inf_true + A_true * (N_million ** (-alpha_true)) + rng.normal(0, 0.01, size=N_million.shape)
L_of_D = L_inf_true + B_true * (D_billion ** (-beta_true)) + rng.normal(0, 0.01, size=D_billion.shape)

print('N sweep losses:', np.round(L_of_N, 4))
print('D sweep losses:', np.round(L_of_D, 4))
print('unit note: N uses params count, D uses token count in isoflops section')

観測点から指数を拾う

ここでは、合成観測値に対して alpha, beta, $L_inf$ を当てにいきます。フィットが合っているか以上に、どの仮定が結果を支配しているかに注目してください。

plt.figure(figsize=(7.2, 3.6))
plt.subplot(1, 2, 1)
plt.plot(N_million, L_of_N, marker='o')
plt.xscale('log')
plt.xlabel('N (million params, log)')
plt.ylabel('validation loss')
plt.title('Loss vs Model Size')

plt.subplot(1, 2, 2)
plt.plot(D_billion, L_of_D, marker='o', color='#d77f00')
plt.xscale('log')
plt.xlabel('D (billion tokens, log)')
plt.ylabel('validation loss')
plt.title('Loss vs Data Size')
plt.tight_layout()
plt.show()

$L_inf$ が未知なので、候補値を走査しながら
log(L - L_inf) と log(x) の一次回帰で指数を推定します。

$L_inf$ は観測損失の最小値より少し小さい値にしかなりえないので、
その近傍を候補として探索します。
ここでの探索範囲は、合成データを扱う教育用の便宜的な設定です。実データでは観測最小値からかなり離れた $L_inf$ が推定されることもあるため、より広い候補範囲や制約付き最適化で floor を推定するのが一般的です。

def fit_power_with_fixed_floor(x, y, floor):
    # y ~= floor + A * x^{-exponent}
    z = y - floor
    if np.any(z <= 0):
        return None

    lx = np.log(x)
    lz = np.log(z)
    # lz = c + m*lx, where m=-exponent
    m, c = np.polyfit(lx, lz, 1)
    pred = floor + np.exp(c) * (x ** m)
    mse = float(np.mean((pred - y) ** 2))
    return {
        'floor': float(floor),
        'A': float(np.exp(c)),
        'exponent': float(-m),
        'mse': mse,
        'pred': pred,
    }


min_obs = float(min(np.min(L_of_N), np.min(L_of_D)))
floor_candidates = np.linspace(min_obs - 0.25, min_obs - 1e-4, 400)

best = None
for floor in floor_candidates:
    fN = fit_power_with_fixed_floor(N_million, L_of_N, floor)
    fD = fit_power_with_fixed_floor(D_billion, L_of_D, floor)
    if fN is None or fD is None:
        continue
    total_mse = fN['mse'] + fD['mse']
    if best is None or total_mse < best['total_mse']:
        best = {
            'L_inf': float(floor),
            'N_fit': fN,
            'D_fit': fD,
            'total_mse': float(total_mse),
        }

fit_joint = best
print('Shared-floor fit summary:')
print('L_inf =', round(fit_joint['L_inf'], 5), 'total_mse =', round(fit_joint['total_mse'], 7))
print('N_fit =', {k: round(v, 5) for k, v in fit_joint['N_fit'].items() if k != 'pred'})
print('D_fit =', {k: round(v, 5) for k, v in fit_joint['D_fit'].items() if k != 'pred'})

同じ計算資源なら、どの配分が得かを見る

スケーリング則が便利なのはここからです。予算を固定したまま、モデルサイズとデータ量のどちらへ振ると損失がより下がるかを見ます。

plt.figure(figsize=(7.2, 3.5))
plt.subplot(1, 2, 1)
plt.scatter(N_million, L_of_N, label='observed')
plt.plot(N_million, fit_joint['N_fit']['pred'], label='power fit', color='#cc3344')
plt.xscale('log')
plt.xlabel('N (M params)')
plt.ylabel('loss')
plt.title(f"alpha≈{fit_joint['N_fit']['exponent']:.3f}")
plt.legend()

plt.subplot(1, 2, 2)
plt.scatter(D_billion, L_of_D, label='observed')
plt.plot(D_billion, fit_joint['D_fit']['pred'], label='power fit', color='#cc3344')
plt.xscale('log')
plt.xlabel('D (B tokens)')
plt.ylabel('loss')
plt.title(f"beta≈{fit_joint['D_fit']['exponent']:.3f}")
plt.legend()
plt.tight_layout()
plt.show()

次に isoflops を考えます。
デコーダ型学習の粗い近似として C ≈ 6ND（N: パラメータ数, D: 学習トークン数）を使います。

D = C/(6N) を L(N, D) に代入すると
f(N) = aN^{-alpha} + b(C/6)^{-beta}N^{beta} になり、
これを df/dN = 0 で解くと N* と D* が得られます。

# 共有L_infで推定した係数を使って、L(N,D)=L_inf+aN^-alpha+bD^-beta を最適化
L_inf = fit_joint['L_inf']
a_fit, alpha = fit_joint['N_fit']['A'], fit_joint['N_fit']['exponent']
b_fit, beta = fit_joint['D_fit']['A'], fit_joint['D_fit']['exponent']

# a_fit,b_fit は N(M params), D(B tokens) の単位で推定されているので
# isoflops (N: params, D: tokens) へ合わせて係数を変換する
# (N/1e6)^-alpha = (1e6^alpha) * N^-alpha
# (D/1e9)^-beta  = (1e9^beta)  * D^-beta
a_raw = a_fit * (1e6 ** alpha)
b_raw = b_fit * (1e9 ** beta)


def optimal_N_D_for_compute(C, a, alpha, b, beta):
    # C = 6ND -> D = C/(6N)
    # minimize f(N)=aN^-alpha + b(C/6)^-beta N^beta
    numer = a * alpha
    denom = b * beta
    N_star = (numer / denom) ** (1.0 / (alpha + beta)) * (C / 6.0) ** (beta / (alpha + beta))
    D_star = C / (6.0 * N_star)
    return N_star, D_star


# C の単位は FLOPs 相当の抽象値（ここでは比較目的）
C_values = np.logspace(18, 22, 9)
N_star = []
D_star = []
for C in C_values:
    n, d = optimal_N_D_for_compute(C, a_raw, alpha, b_raw, beta)
    N_star.append(n)
    D_star.append(d)
N_star = np.array(N_star)
D_star = np.array(D_star)

print('first 4 optimal pairs (N*, D*):')
for i in range(4):
    print(f"C={C_values[i]:.1e} -> N*={N_star[i]/1e6:.2f}M params, D*={D_star[i]/1e9:.2f}B tokens")

配分を誤ると何が起こるかを比べる

後半では、undertrained な大規模モデルと、十分に学習された中規模モデルのような典型的な失敗を並べます。数字の差が、そのまま計画ミスの差として見えるようにします。

plt.figure(figsize=(7.2, 3.6))
plt.plot(C_values, N_star / 1e6, marker='o', label='optimal N* (M params)')
plt.plot(C_values, D_star / 1e9, marker='s', label='optimal D* (B tokens)')
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Compute C (log)')
plt.ylabel('Optimal scale (log)')
plt.title('Isoflops-optimal allocation')
plt.legend()
plt.tight_layout()
plt.show()

実務でありがちな失敗は、計算資源が同じなのに

モデルだけ大きくしてデータ不足（undertrained）
データだけ増やしてモデル不足（underparameterized）

になることです。下で同一 C に対する損失差を比較します。

def approx_loss(N, D, L_inf, a, alpha, b, beta):
    return L_inf + a * (N ** (-alpha)) + b * (D ** (-beta))


ratios = [0.5, 1.0, 2.0]  # Nを最適比の何倍にするか
example_C = 1e20
n_opt, d_opt = optimal_N_D_for_compute(example_C, a_raw, alpha, b_raw, beta)

print(f'compute C={example_C:.1e}')
print(f'optimal N={n_opt/1e6:.2f}M params, D={d_opt/1e9:.2f}B tokens')

for r in ratios:
    n = n_opt * r
    d = example_C / (6.0 * n)
    L = approx_loss(n, d, L_inf, a_raw, alpha, b_raw, beta)
    label = 'optimal' if abs(r - 1.0) < 1e-9 else f'N x {r}'
    print(f'{label:8s}: N={n/1e6:8.2f}M, D={d/1e9:8.2f}B, approx loss={L:.5f}')

スケーリング則は万能ではありません。
データ品質、ドメインミスマッチ、最適化設定、アーキテクチャ変更で指数や下限は変わります。
それでも、実験計画の初期段階で「どこに計算資源を使うか」を決める強力な指針になります。

# 価格の粗い見積もり（仮定値）
train_flops = 3.0e22
hardware_tflops = 250.0   # 1 GPUあたり
num_gpus = 64
utilization = 0.35

seconds = train_flops / (hardware_tflops * 1e12 * num_gpus * utilization)
hours = seconds / 3600

print('Estimated wall-clock hours:', round(hours, 2))

usd_per_gpu_hour = 1.8
cost = hours * num_gpus * usd_per_gpu_hour
print('Estimated training cost (USD, rough):', round(cost, 2))

このノートで押さえたい実務ポイント:

まず小規模スイープで alpha, beta, L_inf を推定する
その推定に基づき isoflops で N と D を配分する
本番では品質劣化要因（データ品質・最適化不安定）を別監視する

この3段階を回すと、計算予算の無駄打ちを減らしやすくなります。