Skip to content
第3回:プラネタリウムの建設 ~極座標へのパラダイムシフト~

第3回:プラネタリウムの建設 ~極座標へのパラダイムシフト~

注意事項

本回で扱う「球面上の学習」には、成立条件がある:

  • 「球面上で計算する」という設計が機能するのは、すべてのベクトルが L2 正規化(ノルム1)されている場合である。ただし、正規化で球面に戻す方法(射影/リトラクション)は、球面制約を満たす実装としては妥当だが、リーマン最適化を厳密に実行していることとは異なる。学習率や運動量の扱いによっては近似的な挙動になる。
  • nGPTはこの設計を徹底するが、標準的なTransformerは必ずしもそうではない。
  • vMF分布の解釈は、埋め込みを方向分布として明示的にモデル化した場合に有効であり、すべての埋め込み空間に自動的に当てはまるわけではない。
  • 「球面上で考えれば常に良い」わけではなく、タスクやデータの性質によっては、ノルムに意味がある場合もある。

導入:平面から球面へ

第2回で、ノルムが「三重の意味」を担っていたことを見た。意味の強度、確信度、スケーリングノイズ——これらが混在し、制御が困難だった。

一つの解決策は、ノルムを固定することである。すべてのベクトルを単位球面上(ノルム1)に制約すれば、ノルムは定数となり、「向き」だけが自由度として残る。

これは単なる正規化テクニックではない。空間そのものを変えるという幾何学的な決断である。ユークリッド空間 Rd から、単位球面 Sd1 へ。平らな世界から、曲がった世界へ。

本回では、この「球面への跳躍」を、nGPTの設計思想、プラネタリウムというメタファー、そしてvon Mises-Fisher分布という数学的道具を通じて理解する。

nGPTの衝撃:設計の徹底

nGPTとは何か

nGPT(Normalized GPT) は、2024年10月にNVIDIAのLoshchilovらが発表したアーキテクチャである(Loshchilov et al., 2024)。その核心的なアイデアは、Transformerのすべての表現を単位球面上に制約することである。

具体的には:

  • 埋め込みベクトル:正規化
  • 注意機構の出力:正規化
  • FFNの出力:正規化
  • 残差接続後:正規化

これにより、同等の精度に達するまでの学習ステップ数を4〜20倍削減したと報告されている(シーケンス長に依存)。

NOTE

Transformerの構成要素について: 上記のリストに登場する「注意機構」「FFN」「残差接続」などが何を指すか分からなくても、本回の議論は追える。これらはTransformerというアーキテクチャの主要な構成要素であり、nGPTはそのすべてに球面制約を課すという設計である。Transformer全体の構造と、各構成要素が幾何学的にどう再解釈されるかについては、第6回の冒頭で整理する

NOTE

「高速化」の意味: nGPTの主張は「学習ステップ数の削減」であり、壁時計時間(実時間)やFLOPsがそのまま4〜20倍改善されるとは限らない。正規化のオーバーヘッドがあるため、1ステップあたりの計算コストは増加する可能性がある。総合的な効率改善の程度は、実装やハードウェアに依存する。

「正規化」と「球面制約」の違い

ここで重要な区別がある。

従来の正規化(例:LayerNorm):

  • 統計量(平均・分散)を正規化: xμσ
  • ベクトルのノルムは必ずしも1にならない
  • 「スケールを安定させる」が目的
  • 後続の線形層(射影行列)で長さが変わりうる

nGPTの球面制約( L2 正規化):

  • 幾何学的な長さを正規化: x/x
  • ベクトルのノルムを厳密に1に固定
  • 「球面上で計算する」が目的
  • 各演算後に再正規化して球面制約を維持

NOTE

LayerNormと球面制約の本質的な違い: LayerNormは「各サンプル内の要素間の統計量」で正規化するのに対し、 L2 正規化は「幾何学的なベクトルの長さ」で正規化する。LayerNormの後に線形層を適用すると、出力の長さは射影行列に依存して変化する。nGPTの論文でも議論されているように、LayerNormは厳密な球面制約を保証しない。

IMPORTANT

設計の徹底が本質: 単に F.normalize を呼ぶだけでは、球面上の学習にはならない。行列積や加算などの演算を行うと、結果は一般に球面から外れる。nGPTは、更新で球面から外れたら正規化で戻すことで、表現が常に球面上にあることを保証する。(多様体最適化の文脈では、こうした「多様体へ戻す写像」を retraction と呼ぶ。)この「球面制約の維持」が、単なる正規化との違いである。

なぜ高速化するのか(仮説)

nGPTが高速化する理由について、いくつかの仮説が考えられる:

仮説説明
勾配の安定化ノルムが固定されると、勾配のスケールも安定する
損失ランドスケープの改善球面上では、局所解の構造が単純になる可能性
学習率の設定が容易ノルムが一定なので、学習率の調整が楽になる
高次元の集中現象との相性下記参照

NOTE

高次元の集中現象(Concentration of Measure)との関連: 第2回で見たように、高次元空間ではガウス分布に従うデータの大半は「原点付近」ではなく「ある半径の殻(シェル)」上に集中する。つまり、高次元データは「自然に球面状に分布しやすい」性質を持つ。nGPTの球面制約は、この高次元の幾何学的性質と相性が良い可能性がある。

CAUTION

仮説段階: 上記は提案されている説明であり、厳密に証明されたメカニズムではない。nGPTの高速化の正確な原因は、さらなる研究が必要である。

歴史的文脈:なぜ2010年代には普及しなかったか

球面制約という発想自体は新しくない。例えば、NormFace(Wang et al., 2017)は顔認証において L2 正規化の有効性を示している。しかし、Transformer全体に球面制約を適用する設計は、2010年代には広く普及しなかった。その理由として、以下が考えられる:

  1. スケーリングと損失設計の非自明さ: NormFaceでも指摘されているように、正規化を入れた学習は単純ではなく、スケールや損失関数の設計に工夫が必要だった(Wang et al., 2017)
  2. 最適化技術との組み合わせ: 球面制約を維持しながら効率的に最適化する実践的なレシピが確立されていなかった
  3. 相対的なオーバーヘッド: 当時のモデル規模では、正規化の計算コストが相対的に大きかった可能性がある

CAUTION

歴史的推測: 上記は「なぜ普及しなかったか」についての推測であり、確定的な因果関係ではない。球面上の最適化自体は、多様体最適化(Absil et al., 2008; Bonnabel, 2013)として理論的には研究されていた。nGPTの貢献は、これをTransformerに適用し、大規模言語モデルで有効性を実証した点にある。

現在は、以下の技術が整っていることが、球面制約の実用化を後押ししている可能性がある:

  • AdamWなどの最適化手法: 重み減衰と適応的学習率の組み合わせ
  • 残差接続: 勾配の流れを安定させる
  • 混合精度学習: 計算効率の向上
  • 大規模GPUメモリ: より大きなバッチサイズが可能に

プラネタリウム・メタファー

球面を「ドーム」として捉える

nGPTの設計を直感的に理解するため、プラネタリウムというメタファーを導入しよう。

プラネタリウムでは、星々がドーム(半球面)に投影される。私たちは地上からドームを見上げ、星の「配置」——星座——を認識する。星までの距離は(投影されているので)問題にならない。重要なのは、星同士の相対的な位置関係(角度) である。

このメタファーを表現学習に適用する:

プラネタリウム表現学習
ドーム単位球面 Sd1
データ点(埋め込みベクトル)
星座意味的なクラスター
星同士の角度類似度
星の明るさ(この比喩では使わない——ノルムを捨てた)

角度が意味を担う

プラネタリウムでは、二つの星の「近さ」は角度で測られる。球面上の角度(測地線距離)は、単位ベクトルの内積と直接対応する:

cosθ=uv(u=v=1)
  • cosθ=1θ=0 ):同じ方向、完全に類似
  • cosθ=0θ=90 ):直交、無関係
  • cosθ=1θ=180 ):反対方向、対立

これがコサイン類似度の幾何学的意味である。

NOTE

単位正規化が前提: コサイン類似度の式 uvuv は、一般のベクトルに対して定義される。しかし、 u=v=1 のとき、これは単に内積 uv になる。nGPTのように球面制約がある場合、コサイン類似度と内積は同一である。

測地線距離

球面上の「最短経路」は測地線(大円の弧)である。二点間の測地線距離は:

dgeo(u,v)=arccos(uv)=θ

これは角度そのものである。第0回で導入した測地線の概念が、ここで具体化される。

von Mises-Fisher分布:球面上の「正規分布」

なぜ球面上の確率分布が必要か

プラネタリウムで「星座」を表現するには、単に点を配置するだけでなく、点の分布を考える必要がある。

ユークリッド空間では、正規分布(ガウス分布)が「最も自然な」分布として使われる。では、球面上では何が対応するのか。

その答えが von Mises-Fisher (vMF) 分布 である。

vMF分布の定義

d 次元単位球面 Sd1 上のvMF分布は、以下の確率密度関数を持つ:

p(xμ,κ)=Cd(κ)exp(κμx)

ここで:

パラメータ意味プラネタリウムでの解釈
μ平均方向(単位ベクトル)星座の「中心」
κ0集中度パラメータ星の「まとまり具合」
Cd(κ)正規化定数(確率の和が1になるよう調整)

正規化定数は:

Cd(κ)=κd/21(2π)d/2Id/21(κ)

ここで Iν は第一種変形ベッセル関数。

集中度 κ の直感

κ は「分布の鋭さ」を制御する:

vMF分布の集中度パラメータ

図: von Mises-Fisher分布における集中度パラメータ κ の効果。 κ=0 では球面上の一様分布となり、 κ が大きくなるほど平均方向 μ への集中度が増す。正規分布の分散 σ2 に対応するのは 1/κ である。

κ の値分布の形状プラネタリウムでの解釈
κ=0球面上の一様分布星がドーム全体にランダムに散らばる
κ 小(例:1〜5)緩やかに集中星座がぼんやり見える
κ 中(例:10〜50)明確に集中星座がはっきり見える
κ一点に集中(デルタ関数に近づく)星が一点に凝縮

vMFと正規分布の対応

vMF分布は、多くの点で正規分布と類似している:

正規分布( RdvMF分布( Sd1
平均 μ平均方向 μ
分散 σ2 (または共分散)集中度の逆数 1/κ
等方的な広がり球面上での等方的な広がり
指数関数的な減衰指数関数的な減衰(角度に対して)

NOTE

ガウスから誘導される「方向分布」としての解釈: 等方的なガウス zN(m,σ2I) を考え、方向ベクトル u=zzSd1 を取ると、 u は vMF 分布に従う(厳密には「 z=1 に条件付ける」という表現よりも、「正規分布のサンプルの方向成分を見ている」と捉えるのが安全)。 このとき、平均方向は μ=m|m 、集中度は κ=mσ2 に対応する。 この関係から「球面上の正規分布」と呼ばれることがあるが、数学的には異なる性質も持つ。

NOTE

最大エントロピー原理: vMF分布は、「平均方向 E[x] が特定の値を取る」という制約の下で、球面上のエントロピーを最大化する分布として特徴づけられる。この意味で「最も仮定の少ない方向分布」と言えるが、制約なしの最大エントロピー分布(一様分布)とは異なる。

Softmaxとの関係

興味深いことに、vMF分布の形式はSoftmax関数と形式的に対応している。

前提条件: 以下の対応が成り立つのは、次の条件が満たされる場合である:

  • クラス k に対応する代表ベクトル μk単位ノルム
  • 入力の埋め込み x単位ノルム
  • ロジットが内積(またはスケール付き内積)として計算される

この条件下で、Softmax分類の形式は:

p(kx)=exp(κμkx)jexp(κμjx)

これは、温度 T=1/κ のSoftmax関数に対応する。

CAUTION

「対応」であって「同一」ではない: vMF分布は連続的な方向(球面全体)上の確率分布であるのに対し、Softmax分類は離散的なカテゴリ上の確率分布である。両者は「指数関数的に内積が入る」という形式が対応しているだけで、確率モデルとして同一視することはできない。

CAUTION

前提が崩れる場合: 標準的なTransformerの分類ヘッドでは、出力層の重み行列や入力埋め込みが正規化されていないことが多い。この場合、ロジットは内積ではなく一般の線形変換となり、vMFとの形式的対応も成り立たない。nGPTのように球面制約を徹底する場合に、この対応が意味を持つ。

IMPORTANT

第4回への伏線: この対応は、Softmaxを情報幾何学の観点から再解釈する際の出発点となる。温度パラメータ T (または κ )が、確率分布空間でどのような役割を果たすかを、次回詳しく見る。

球面上の学習:幾何学的な意味

勾配の方向

球面上で最適化を行うとき、勾配は球面に接する方向でなければならない。球面から離れる方向(半径方向)への勾配成分は、制約を破るからである。

具体的には、点 xSd1 における接空間は:

TxSd1={vRd:vx=0}

すなわち、 x と直交するベクトルの集合。

ユークリッド空間での勾配 f(x) を接空間に射影するには:

gradSf(x)=f(x)(f(x)x)x

これがリーマン勾配(球面上の勾配)である。

正規化による近似

実際の実装では、リーマン勾配を明示的に計算する代わりに、以下の近似がよく使われる:

  1. 通常のユークリッド勾配で更新
  2. 結果を球面に射影(正規化)
python
# 疑似コード
# x = x - lr * grad  # ユークリッド空間での更新
# x = x / x.norm()  # 球面への射影

これは「射影勾配法」の一種であり、学習率が十分小さければ、リーマン勾配降下法に近い挙動を示す。

CAUTION

近似の限界: 学習率が大きい場合、この近似は不正確になる。また、運動量を使う場合、運動量ベクトルも接空間上にあるべきだが、この近似ではそれが保証されない。より厳密なリーマン最適化手法については、Bonnabel (2013) などを参照。

球面設計の限界と注意点

ノルムに意味がある場合

球面制約は万能ではない。ノルムが意味を持つタスクでは、ノルムを捨てることで情報が失われる:

タスクノルムの意味球面制約の適否
分類(表現の比較)不要なことが多い✓ 適している
回帰(連続値の予測)出力の大きさが重要✗ 適さない
不確実性推定ノルムが確信度を表す場合△ タスク依存
生成モデル分布のスケールが重要△ 設計次第

球面と超球面の混同

「球面」という言葉には注意が必要:

用語数学的定義
1次元球面 S1R2 内の円単位円
2次元球面 S2R3 内の球面地球表面
(d1) 次元球面 Sd1Rd 内の超球面高次元の単位球面

768次元の埋め込みを正規化すると、それは S767 上の点になる。

曲率の影響

球面は正の曲率を持つ多様体である。第0回で触れたように、正曲率空間では三角形の内角の和が180°を超える。

この曲率は、学習ダイナミクスに影響を与える可能性がある。例えば、球面上では「遠い二点を直線で結ぶ」ことができない(測地線は曲がる)。

NOTE

双曲空間との対比: 第12回では、負の曲率を持つ双曲空間を扱う。階層構造の表現には、球面より双曲空間が適していることがある。空間の曲率は、表現したい構造に応じて選ぶべきである。

次回予告

第4回「分類の再統一 I」では、Softmaxを情報幾何学の観点から再解釈する。

vMF分布とSoftmaxの関係を深掘りし、「温度」パラメータの幾何学的意味を理解する。さらに、自然勾配がなぜ有効なのかを、確率分布空間の構造から説明する。

TIP

Transformerに馴染みのない読者へ: 本回ではnGPTの設計思想を通じて「球面制約」の意味を理解した。nGPTが改良の対象としているTransformerそのものの全体像——どんな構成要素があり、それぞれが何をしているのか——については、第6回「Transformerという測量士」の冒頭で整理する。第4回・第5回はSoftmaxとマージンという、Transformerに限定されない概念を扱うので、Transformerの詳細を知らなくても読み進められる。

実装ノート:正規化の実践

NOTE

以下のコードは PyTorch >= 1.9 を前提とする。

基本的な正規化

コード例: 03_l2_normalize.py
python
import torch
import torch.nn.functional as F


def l2_normalize(x, dim=-1, eps=1e-12):
    """L2正規化(球面への射影)

    Args:
        x: 入力テンソル
        dim: 正規化する次元
        eps: 数値安定性のための小さな値

    Returns:
        単位ノルムに正規化されたテンソル
    """
    return F.normalize(x, p=2, dim=dim, eps=eps)


# 使用例
embeddings = torch.randn(32, 768)  # バッチサイズ32、次元768
normalized = l2_normalize(embeddings)

# 確認:ノルムが1になっている
print(f"ノルムの平均: {normalized.norm(dim=-1).mean():.6f}")  # ≈ 1.0
print(f"ノルムの標準偏差: {normalized.norm(dim=-1).std():.6f}")  # ≈ 0.0

数値安定性の問題

正規化には数値安定性の問題がある。特に、混合精度学習(float16)ではアンダーフローが起きやすい:

コード例: 03_fp16_instability_demo.py
python
import torch
import torch.nn.functional as F


def demonstrate_instability():
    """float16での数値不安定性の例"""

    # float16では非常に小さい値が0に丸められやすい(アンダーフロー)
    # min positive normal ≈ 6e-5 なので、それより十分小さい値は 0 になりやすい
    very_small = torch.finfo(torch.float16).tiny / 1024  # float32では非ゼロ
    small_fp16 = torch.full((3,), very_small, dtype=torch.float32).half()  # → 0 になりうる

    print(f"入力 (float16): {small_fp16}")
    print(f"ノルム: {small_fp16.norm()}")  # 非常に小さい、または0

    # eps なしで正規化すると不安定(ノルムが 0 に丸められると NaN/Inf)
    try:
        normalized_unsafe = small_fp16 / small_fp16.norm()
        print(f"Unsafe (epsなし): {normalized_unsafe}")
    except Exception as e:
        print(f"Unsafe normalization failed: {e}")

    # 実務では「ゼロベクトルが混ざる」ケースも多い(マスクや初期化など)
    zero_fp16 = torch.zeros(3, dtype=torch.float16)
    print(f"ゼロ入力 (float16): {zero_fp16}, ノルム: {zero_fp16.norm()}")
    print(f"ゼロをepsなしで割る: {zero_fp16 / zero_fp16.norm()}")  # NaN が出る

    # 安全策:ノルム計算をfloat32で行い、epsで保護して正規化
    normalized_safe = F.normalize(small_fp16.float(), eps=1e-12).half()
    print(f"Safe (eps使用, normはfloat32): {normalized_safe}")

    # 勾配計算時にゼロ近傍で問題が起きる例
    x = torch.tensor([1e-7, 1e-7, 1e-7], requires_grad=True)
    y = x / (x.norm() + 1e-12)  # eps で保護
    y.sum().backward()
    print(f"勾配 (eps保護あり): {x.grad}")


demonstrate_instability()

CAUTION

混合精度学習での注意: 大規模モデルでは計算効率のためにfloat16やbfloat16が使われることが多い。正規化の際は、ノルム計算をfloat32で行うか、十分なepsを設定することが重要。

nGPT風の設計:球面制約の維持

コード例: 03_normalized_linear.py
python
import torch
import torch.nn as nn
import torch.nn.functional as F


class NormalizedLinear(nn.Module):
    """正規化を含む線形層(nGPT風)"""

    def __init__(self, in_features, out_features, eps=1e-12):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features, bias=False)
        self.eps = eps

        # 重みも正規化して初期化
        with torch.no_grad():
            self.linear.weight.data = F.normalize(self.linear.weight.data, dim=1, eps=eps)

    def forward(self, x):
        # 入力は正規化済みと仮定
        out = self.linear(x)
        # 出力を正規化
        return F.normalize(out, dim=-1, eps=self.eps)


class NormalizedResidual(nn.Module):
    """正規化を含む残差接続(nGPT風)"""

    def __init__(self, module, alpha=0.1, eps=1e-12):
        super().__init__()
        self.module = module
        self.alpha = alpha  # 残差のスケール
        self.eps = eps

    def forward(self, x):
        # 残差を加算
        out = x + self.alpha * self.module(x)
        # 正規化して球面に戻す
        return F.normalize(out, dim=-1, eps=self.eps)

IMPORTANT

概念と実装のギャップ: 上記のコードは「nGPT風」の設計を示すものであり、nGPTの正確な実装ではない。nGPTの詳細は原論文(Loshchilov et al., 2024)を参照されたい。

vMF分布の可視化

コード例: 03_vmf_distribution_plot.py
python
import matplotlib.pyplot as plt
import numpy as np
from scipy.special import iv  # ベッセル関数


def vmf_pdf_2d(theta, kappa):
    """2次元(円上)のvMF分布の確率密度

    平均方向を θ=0 として、角度 θ での密度を返す
    """
    # 2次元の場合、正規化定数は 1/(2π I_0(κ))
    normalization = 1 / (2 * np.pi * iv(0, kappa))
    return normalization * np.exp(kappa * np.cos(theta))


# 異なる κ での分布を可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 左:確率密度のプロット
theta = np.linspace(-np.pi, np.pi, 1000)
kappas = [0, 1, 5, 10, 50]

ax = axes[0]
for kappa in kappas:
    if kappa == 0:
        # κ=0 は一様分布
        pdf = np.ones_like(theta) / (2 * np.pi)
    else:
        pdf = vmf_pdf_2d(theta, kappa)
    ax.plot(theta, pdf, label=f"κ = {kappa}")

ax.set_xlabel("Angle θ (radians)")
ax.set_ylabel("Probability Density")
ax.set_title("von Mises-Fisher Distribution (2D)")
ax.legend()
ax.set_xlim(-np.pi, np.pi)

# 右:極座標でのプロット(円上の分布)
ax = axes[1]
for kappa in kappas:
    if kappa == 0:
        r = np.ones_like(theta) / (2 * np.pi)
    else:
        r = vmf_pdf_2d(theta, kappa)
    # 極座標での x, y
    x = r * np.cos(theta)
    y = r * np.sin(theta)
    ax.plot(x, y, label=f"κ = {kappa}")

ax.set_aspect("equal")
ax.set_title("vMF Distribution on Circle")
ax.legend()

plt.tight_layout()
plt.show()

参考文献

nGPTと球面上の学習

  • Loshchilov, I., Hsieh, C.-P., Sun, S., & Ginsburg, B. (2024). nGPT: Normalized Transformer with Representation Learning on the Hypersphere. arXiv:2410.01131.
    • nGPTの原論文。Transformerのすべての表現を単位球面上に制約する設計。
  • Wang, F., Xiang, X., Cheng, J., & Yuille, A. L. (2017). NormFace: L2 Hypersphere Embedding for Face Verification. ACM MM 2017. arXiv: arXiv:1704.06369.
    • 顔認証における球面埋め込みの先駆的研究。

von Mises-Fisher分布

  • Mardia, K. V., & Jupp, P. E. (2000). Directional Statistics. Wiley.
    • 方向統計学の標準的教科書。vMF分布を含む球面上の確率分布を体系的に扱う。
  • Banerjee, A., Dhillon, I. S., Ghosh, J., & Sra, S. (2005). Clustering on the Unit Hypersphere using von Mises-Fisher Distributions. JMLR, 6, 1345–1382.
    • vMF分布を用いたクラスタリング。機械学習への応用。

リーマン最適化

  • Bonnabel, S. (2013). Stochastic Gradient Descent on Riemannian Manifolds. IEEE Transactions on Automatic Control, 58(9), 2217–2229.
    • 多様体上の確率的勾配降下法の理論的基礎。
  • Absil, P.-A., Mahony, R., & Sepulchre, R. (2008). Optimization Algorithms on Matrix Manifolds. Princeton University Press.
    • 行列多様体(Stiefel多様体、Grassmann多様体など)上の最適化。球面も特殊ケースとして含む。