第3回:プラネタリウムの建設 ~極座標へのパラダイムシフト~
注意事項
本回で扱う「球面上の学習」には、成立条件がある:
- 「球面上で計算する」という設計が機能するのは、すべてのベクトルが
正規化(ノルム1)されている場合である。ただし、正規化で球面に戻す方法(射影/リトラクション)は、球面制約を満たす実装としては妥当だが、リーマン最適化を厳密に実行していることとは異なる。学習率や運動量の扱いによっては近似的な挙動になる。 - nGPTはこの設計を徹底するが、標準的なTransformerは必ずしもそうではない。
- vMF分布の解釈は、埋め込みを方向分布として明示的にモデル化した場合に有効であり、すべての埋め込み空間に自動的に当てはまるわけではない。
- 「球面上で考えれば常に良い」わけではなく、タスクやデータの性質によっては、ノルムに意味がある場合もある。
導入:平面から球面へ
第2回で、ノルムが「三重の意味」を担っていたことを見た。意味の強度、確信度、スケーリングノイズ——これらが混在し、制御が困難だった。
一つの解決策は、ノルムを固定することである。すべてのベクトルを単位球面上(ノルム1)に制約すれば、ノルムは定数となり、「向き」だけが自由度として残る。
これは単なる正規化テクニックではない。空間そのものを変えるという幾何学的な決断である。ユークリッド空間
本回では、この「球面への跳躍」を、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):
- 統計量(平均・分散)を正規化:
- ベクトルのノルムは必ずしも1にならない
- 「スケールを安定させる」が目的
- 後続の線形層(射影行列)で長さが変わりうる
nGPTの球面制約(
- 幾何学的な長さを正規化:
- ベクトルのノルムを厳密に1に固定
- 「球面上で計算する」が目的
- 各演算後に再正規化して球面制約を維持
NOTE
LayerNormと球面制約の本質的な違い: LayerNormは「各サンプル内の要素間の統計量」で正規化するのに対し、
IMPORTANT
設計の徹底が本質: 単に F.normalize を呼ぶだけでは、球面上の学習にはならない。行列積や加算などの演算を行うと、結果は一般に球面から外れる。nGPTは、更新で球面から外れたら正規化で戻すことで、表現が常に球面上にあることを保証する。(多様体最適化の文脈では、こうした「多様体へ戻す写像」を retraction と呼ぶ。)この「球面制約の維持」が、単なる正規化との違いである。
なぜ高速化するのか(仮説)
nGPTが高速化する理由について、いくつかの仮説が考えられる:
| 仮説 | 説明 |
|---|---|
| 勾配の安定化 | ノルムが固定されると、勾配のスケールも安定する |
| 損失ランドスケープの改善 | 球面上では、局所解の構造が単純になる可能性 |
| 学習率の設定が容易 | ノルムが一定なので、学習率の調整が楽になる |
| 高次元の集中現象との相性 | 下記参照 |
NOTE
高次元の集中現象(Concentration of Measure)との関連: 第2回で見たように、高次元空間ではガウス分布に従うデータの大半は「原点付近」ではなく「ある半径の殻(シェル)」上に集中する。つまり、高次元データは「自然に球面状に分布しやすい」性質を持つ。nGPTの球面制約は、この高次元の幾何学的性質と相性が良い可能性がある。
CAUTION
仮説段階: 上記は提案されている説明であり、厳密に証明されたメカニズムではない。nGPTの高速化の正確な原因は、さらなる研究が必要である。
歴史的文脈:なぜ2010年代には普及しなかったか
球面制約という発想自体は新しくない。例えば、NormFace(Wang et al., 2017)は顔認証において
- スケーリングと損失設計の非自明さ: NormFaceでも指摘されているように、正規化を入れた学習は単純ではなく、スケールや損失関数の設計に工夫が必要だった(Wang et al., 2017)
- 最適化技術との組み合わせ: 球面制約を維持しながら効率的に最適化する実践的なレシピが確立されていなかった
- 相対的なオーバーヘッド: 当時のモデル規模では、正規化の計算コストが相対的に大きかった可能性がある
CAUTION
歴史的推測: 上記は「なぜ普及しなかったか」についての推測であり、確定的な因果関係ではない。球面上の最適化自体は、多様体最適化(Absil et al., 2008; Bonnabel, 2013)として理論的には研究されていた。nGPTの貢献は、これをTransformerに適用し、大規模言語モデルで有効性を実証した点にある。
現在は、以下の技術が整っていることが、球面制約の実用化を後押ししている可能性がある:
- AdamWなどの最適化手法: 重み減衰と適応的学習率の組み合わせ
- 残差接続: 勾配の流れを安定させる
- 混合精度学習: 計算効率の向上
- 大規模GPUメモリ: より大きなバッチサイズが可能に
プラネタリウム・メタファー
球面を「ドーム」として捉える
nGPTの設計を直感的に理解するため、プラネタリウムというメタファーを導入しよう。
プラネタリウムでは、星々がドーム(半球面)に投影される。私たちは地上からドームを見上げ、星の「配置」——星座——を認識する。星までの距離は(投影されているので)問題にならない。重要なのは、星同士の相対的な位置関係(角度) である。
このメタファーを表現学習に適用する:
| プラネタリウム | 表現学習 |
|---|---|
| ドーム | 単位球面 |
| 星 | データ点(埋め込みベクトル) |
| 星座 | 意味的なクラスター |
| 星同士の角度 | 類似度 |
| 星の明るさ | (この比喩では使わない——ノルムを捨てた) |
角度が意味を担う
プラネタリウムでは、二つの星の「近さ」は角度で測られる。球面上の角度(測地線距離)は、単位ベクトルの内積と直接対応する:
( ):同じ方向、完全に類似 ( ):直交、無関係 ( ):反対方向、対立
これがコサイン類似度の幾何学的意味である。
NOTE
単位正規化が前提: コサイン類似度の式
測地線距離
球面上の「最短経路」は測地線(大円の弧)である。二点間の測地線距離は:
これは角度そのものである。第0回で導入した測地線の概念が、ここで具体化される。
von Mises-Fisher分布:球面上の「正規分布」
なぜ球面上の確率分布が必要か
プラネタリウムで「星座」を表現するには、単に点を配置するだけでなく、点の分布を考える必要がある。
ユークリッド空間では、正規分布(ガウス分布)が「最も自然な」分布として使われる。では、球面上では何が対応するのか。
その答えが von Mises-Fisher (vMF) 分布 である。
vMF分布の定義
ここで:
| パラメータ | 意味 | プラネタリウムでの解釈 |
|---|---|---|
| 平均方向(単位ベクトル) | 星座の「中心」 | |
| 集中度パラメータ | 星の「まとまり具合」 | |
| 正規化定数 | (確率の和が1になるよう調整) |
正規化定数は:
ここで
集中度 の直感
図: von Mises-Fisher分布における集中度パラメータ
| 分布の形状 | プラネタリウムでの解釈 | |
|---|---|---|
| 球面上の一様分布 | 星がドーム全体にランダムに散らばる | |
| 緩やかに集中 | 星座がぼんやり見える | |
| 明確に集中 | 星座がはっきり見える | |
| 一点に集中(デルタ関数に近づく) | 星が一点に凝縮 |
vMFと正規分布の対応
vMF分布は、多くの点で正規分布と類似している:
| 正規分布( | vMF分布( |
|---|---|
| 平均 | 平均方向 |
| 分散 | 集中度の逆数 |
| 等方的な広がり | 球面上での等方的な広がり |
| 指数関数的な減衰 | 指数関数的な減衰(角度に対して) |
NOTE
ガウスから誘導される「方向分布」としての解釈: 等方的なガウス
NOTE
最大エントロピー原理: vMF分布は、「平均方向
Softmaxとの関係
興味深いことに、vMF分布の形式はSoftmax関数と形式的に対応している。
前提条件: 以下の対応が成り立つのは、次の条件が満たされる場合である:
- クラス
に対応する代表ベクトル が単位ノルム - 入力の埋め込み
も単位ノルム - ロジットが内積(またはスケール付き内積)として計算される
この条件下で、Softmax分類の形式は:
これは、温度
CAUTION
「対応」であって「同一」ではない: vMF分布は連続的な方向(球面全体)上の確率分布であるのに対し、Softmax分類は離散的なカテゴリ上の確率分布である。両者は「指数関数的に内積が入る」という形式が対応しているだけで、確率モデルとして同一視することはできない。
CAUTION
前提が崩れる場合: 標準的なTransformerの分類ヘッドでは、出力層の重み行列や入力埋め込みが正規化されていないことが多い。この場合、ロジットは内積ではなく一般の線形変換となり、vMFとの形式的対応も成り立たない。nGPTのように球面制約を徹底する場合に、この対応が意味を持つ。
IMPORTANT
第4回への伏線: この対応は、Softmaxを情報幾何学の観点から再解釈する際の出発点となる。温度パラメータ
球面上の学習:幾何学的な意味
勾配の方向
球面上で最適化を行うとき、勾配は球面に接する方向でなければならない。球面から離れる方向(半径方向)への勾配成分は、制約を破るからである。
具体的には、点
すなわち、
ユークリッド空間での勾配
これがリーマン勾配(球面上の勾配)である。
正規化による近似
実際の実装では、リーマン勾配を明示的に計算する代わりに、以下の近似がよく使われる:
- 通常のユークリッド勾配で更新
- 結果を球面に射影(正規化)
# 疑似コード
# x = x - lr * grad # ユークリッド空間での更新
# x = x / x.norm() # 球面への射影これは「射影勾配法」の一種であり、学習率が十分小さければ、リーマン勾配降下法に近い挙動を示す。
CAUTION
近似の限界: 学習率が大きい場合、この近似は不正確になる。また、運動量を使う場合、運動量ベクトルも接空間上にあるべきだが、この近似ではそれが保証されない。より厳密なリーマン最適化手法については、Bonnabel (2013) などを参照。
球面設計の限界と注意点
ノルムに意味がある場合
球面制約は万能ではない。ノルムが意味を持つタスクでは、ノルムを捨てることで情報が失われる:
| タスク | ノルムの意味 | 球面制約の適否 |
|---|---|---|
| 分類(表現の比較) | 不要なことが多い | ✓ 適している |
| 回帰(連続値の予測) | 出力の大きさが重要 | ✗ 適さない |
| 不確実性推定 | ノルムが確信度を表す場合 | △ タスク依存 |
| 生成モデル | 分布のスケールが重要 | △ 設計次第 |
球面と超球面の混同
「球面」という言葉には注意が必要:
| 用語 | 数学的定義 | 例 |
|---|---|---|
| 1次元球面 | 単位円 | |
| 2次元球面 | 地球表面 | |
| 高次元の単位球面 |
768次元の埋め込みを正規化すると、それは
曲率の影響
球面は正の曲率を持つ多様体である。第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
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
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
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
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多様体など)上の最適化。球面も特殊ケースとして含む。