第3回: プラネタリウムの建設: 球面へのパラダイムシフト

TIP

読了後に、理解度チェックで要点を確認できる。 このページで登場した用語は、キーワード集でも確認できる。

注意事項

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

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

導入：平面から球面へ

第2回で、ノルムが意味の強度、確信度、スケーリングノイズと、「三重の意味」を担っていたことを見た。 3つの意味を後から仕分けようとする限り、この混線はおそらく解けない。

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

これは単なる正規化テクニックではない。空間そのものを変えるという幾何学的な決断である。ユークリッド空間 ${\mathbb{R}}^d$ から、単位球面 ${\mathbb{S}}^{d-1}$ へ。平らな世界から、曲がった世界へ。

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

nGPTの衝撃：設計の徹底

nGPTとは何か

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

具体的には：

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

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

NOTE

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

NOTE

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

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

従来の正規化（例：LayerNorm）：

• 統計量（平均と分散）を正規化： $\frac{x-\mu }{\sigma }$
• ベクトルのノルムは必ずしも1にならない
• 「スケールを安定させる」が目的
• 後続の線形層（射影行列）で長さが変わりうる

nGPTの球面制約（ $L_2$ 正規化）：

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

NOTE

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

IMPORTANT

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

なぜ高速化するのか（仮説）

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

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

NOTE

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

CAUTION

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

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

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

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

CAUTION

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

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

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

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

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

数式だけでは、球面制約の「感じ」はなかなか掴めない。nGPTの設計を直感的に理解するため、プラネタリウムというメタファーを導入しよう。

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

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

プラネタリウム	表現学習
ドーム	単位球面 ${\mathbb{S}}^{d-1}$
星	データ点（埋め込みベクトル）
星座	意味的なクラスター
星同士の角度	類似度
星の明るさ	（この比喩では使わない。ノルムを捨てた）

この対応表に、本講義の核心が凝縮されている。星は、広大な空間に離散的に配置された「意味の所在」である。その星がどこにあるか（方向、配置） だけを見て、どれだけ明るいか（大きさ、ノルム） は見ない。星の明るさをあえて捨てる。これが第2回で見た「ノルムの呪い」への回答である。意味は、星の配置（角度）だけが担う。

CAUTION

プラネタリウムの賞味期限： このメタファーは、空間全体が均一な曲率を持つ球面であることを前提にしている。第12回で扱う双曲空間のように曲率が場所や符号によって異なる一般の多様体では、「ドームに星を貼る」という単純な描像は成り立たない。プラネタリウムは球面という特殊ケースの直感であり、任意の多様体に拡張できるものではない。

NOTE

第1回の「星空」との違い： 第1回でも「星＝単語、星座＝トピック」という比喩を使ったが、あれは離散的な星空だった。星はあっても、それらを浮かべる連続的な空間がなかった。本回で単位球面という連続的な夜空を得て、星は初めて「方向で意味を持つ点」になる。以降の章でも「星」は登場するが、その本質は一貫して「離散的な意味の所在を、配置（方向）で意味づける」ことにある。

角度が意味を担う

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

$$\cos \theta ={\mathbf{u}}^{\mathrm{\top }}\mathbf{v}(\parallel \mathbf{u}\parallel =\parallel \mathbf{v}\parallel =1)$$

• $\cos \theta =1$ （ $\theta =0^{\circ }$ ）：同じ方向、完全に類似
• $\cos \theta =0$ （ $\theta ={90}^{\circ }$ ）：直交、無関係
• $\cos \theta =-1$ （ $\theta ={180}^{\circ }$ ）：反対方向、対立

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

NOTE

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

角度はスカラーだが、分類は1個では終わらない

ここまでの $\theta$ は、二つの星 $\mathbf{u},\mathbf{v}$ の間の角度、すなわち1個のスカラーである。「角度で分類する」という言い方から、「入力を1個の角度に変換し、その数直線上の位置でラベルを決める」という1次元のイメージを持つと、「たった1次元にすべてのクラスを詰め込むのは窮屈ではないか」という誤解が生じやすい。

画像分類を例にすると、処理は「画像 → 特徴抽出 → $d$ 次元の特徴ベクトル $\mathbf{x}$ → ラベル」という流れである。ただし、特徴ベクトルからラベルへ進むとき、画像を1個の角度に変換しているわけではない。 $\mathbf{x}$ を $K$ 個あるクラス代表ベクトル ${\mathbf{w}}_1,\dots ,{\mathbf{w}}_K$ のそれぞれと比較し、

$${\theta }_1=\mathrm{\angle }(\mathbf{x},{\mathbf{w}}_1),{\theta }_2=\mathrm{\angle }(\mathbf{x},{\mathbf{w}}_2),\dots ,{\theta }_K=\mathrm{\angle }(\mathbf{x},{\mathbf{w}}_K)$$

という $K$ 個の角度を同時に計算する。正規化された角度ベース分類器で、推論時に最も近いクラスを選ぶ場合は、この中で最も角度が小さい（＝ $\cos {\theta }_k$ が最大の）クラスがラベルになる。

• 角度 ${\theta }_k$ 単体は、確かに1回の比較につき1個のスカラーである。
• しかし「分類」という営みは、この比較を $K$ 個のクラス代表に対して並列に行い、その中から最小値を選ぶ作業である。
• $\mathbf{x}$ 自身は $d$ 次元の球面 ${\mathbb{S}}^{d-1}$ 上の1点であり、正規化された ${\mathbf{w}}_k$ も同じ球面上にある。角度は、この高次元の配置を「2点間の関係」として覗き見るための1次元の窓に過ぎない。

nGPTは画像分類モデルではなく、文脈から次のトークンを予測するモデルである。ここでは、画像から得た特徴ベクトル $\mathbf{x}$ に相当するものが文脈から得られた表現ベクトルであり、クラス代表ベクトル ${\mathbf{w}}_k$ に相当するものが各候補トークンの出力埋め込みである。したがって、画像分類と次トークン予測は対象こそ異なるが、「高次元の表現ベクトルを各候補と比較し、候補ごとのスコアを得る」という幾何学は共通している。

この「1対1の角度」から「1対 $K$ の並列比較」への切り替えを、第4回でロジット $z_k=\cos {\theta }_k$ として定式化する。

測地線距離

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

$$d_{\text{geo}}(\mathbf{u},\mathbf{v})=\arccos ({\mathbf{u}}^{\mathrm{\top }}\mathbf{v})=\theta$$

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

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

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

プラネタリウムで「星座」を表現するには、単に点を配置するだけでなく、点の分布を考える必要がある。一つの星座を成す星々は、中心からどれだけばらついているのか、という問いだ。

ユークリッド空間では、正規分布（ガウス分布）が広く使われる。 では、球面上では何が対応するのか。 正規分布に従うベクトルを原点から球面へ射影すると、射影正規分布（projected normal distribution） が得られ、これは球面データを表す正当な選択肢である。 実際、中心が原点にある等方的な正規分布を射影した場合は、球面上の一様分布になる。 ただし、一般の射影正規分布は指数型分布族の単純な形を持たず、平均方向と集中度を明示的に分離したい場面では扱いが複雑になる。 そこで本回では、平均方向 $\mathit{\mu }$ と集中度 $\kappa$ を直接指定でき、指数型分布族として解析しやすい分布を選ぶ。

その答えが von Mises-Fisher（vMF）分布 である（Mardia & Jupp, 2000; Banerjee et al., 2005）。

vMF分布の定義

$d$ 次元単位球面 ${\mathbb{S}}^{d-1}$ 上のvMF分布は、以下の確率密度関数を持つ：

$$p(\mathbf{x}\mid \mathit{\mu },\kappa )=C_d(\kappa )\exp (\kappa {\mathit{\mu }}^{\mathrm{\top }}\mathbf{x})$$

ここで：

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

正規化定数は：

$$C_d(\kappa )=\frac{{\kappa }^{d/2-1}}{(2\pi {)}^{d/2}{I}_{d/2-1}(\kappa )}$$

ここで $I_{\nu }$ は第一種変形ベッセル関数。

集中度 $\kappa$ の直感

$\kappa$ は「分布の鋭さ」を制御する：

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

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

vMFと正規分布の対応

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

正規分布（ ${\mathbb{R}}^d$ ）	vMF分布（ ${\mathbb{S}}^{d-1}$ ）
平均 $\mathit{\mu }$	平均方向 $\mathit{\mu }$
分散 ${\sigma }^2$ （または共分散）	集中度の逆数 $1/\kappa$
等方的な広がり	球面上での等方的な広がり
指数関数的な減衰	指数関数的な減衰（角度に対して）

NOTE

ガウスから誘導される「方向分布」としての解釈： 等方的なガウス $\mathbf{z}\sim N(\mathbf{m},{\sigma }^{2}I)$ を考え、方向ベクトル $\mathbf{u}=\frac{\mathbf{z}}{\parallel \mathbf{z}\parallel }\in {\mathbb{S}}^{d-1}$ を取ると、 $\mathbf{u}$ は厳密には projected normal distribution に従う。ただし、大 $\kappa$ かつ高次元の条件下では vMF 分布で近似できる。 このとき、平均方向は $\mathit{\mu }=\frac{\mathbf{m}}{\parallel \mathbf{m}\parallel }$ 、集中度は $\kappa =\frac{\parallel \mathbf{m}\parallel }{{\sigma }^2}$ に対応する（ $\parallel \mathbf{z}\parallel =1$ に条件付けた密度として見た場合）。 この関係から「球面上の正規分布」と呼ばれることがあるが、数学的には異なる性質も持つ。

NOTE

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

Softmaxとの関係

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

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

• クラス $k$ に対応する代表ベクトル ${\mathit{\mu }}_k$ が単位ノルム
• 入力の埋め込み $\mathbf{x}$ も単位ノルム
• ロジットが内積（またはスケール付き内積）として計算される

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

$$p(k\mid \mathbf{x})=\frac{\exp (\kappa {\mathit{\mu }}_k^{\mathrm{\top }}\mathbf{x})}{\sum _j\exp (\kappa {\mathit{\mu }}_j^{\mathrm{\top }}\mathbf{x})}$$

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

CAUTION

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

CAUTION

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

IMPORTANT

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

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

勾配の方向

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

具体的には、点 $\mathbf{x}\in {\mathbb{S}}^{d-1}$ における接空間は：

$$T_{\mathbf{x}}{\mathbb{S}}^{d-1}=\{\mathbf{v}\in {\mathbb{R}}^d:{\mathbf{v}}^{\mathrm{\top }}\mathbf{x}=0\}$$

すなわち、 $\mathbf{x}$ と直交するベクトルの集合。

ユークリッド空間での勾配 $\mathrm{\nabla }f(\mathbf{x})$ を接空間に射影するには：

$${\text{grad}}_{S}f(\mathbf{x})=\mathrm{\nabla }f(\mathbf{x})-(\mathrm{\nabla }f(\mathbf{x}{)}^{\mathrm{\top }}\mathbf{x})\mathbf{x}$$

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

正規化による近似

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

1. 通常のユークリッド勾配で更新
2. 結果を球面に射影（正規化）

# 疑似コード
# x = x - lr * grad  # ユークリッド空間での更新
# x = x / x.norm()  # 球面への射影

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

CAUTION

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

球面設計の限界と注意点

ノルムに意味がある場合

ここまで球面制約の効能を見てきた。しかし万能ではない。ノルムが意味を持つタスクでは、ノルムを捨てることで情報が失われる：

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

次元表記の混同に注意

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

用語	数学的定義	例
1次元球面 ${\mathbb{S}}^1$	${\mathbb{R}}^2$ 内の円	単位円
2次元球面 ${\mathbb{S}}^2$	${\mathbb{R}}^3$ 内の球面	地球表面
$(d-1)$ 次元球面 ${\mathbb{S}}^{d-1}$	${\mathbb{R}}^d$ 内の球面	高次元の単位球面

768次元の埋め込みを正規化すると、それは ${\mathbb{S}}^{767}$ 上の点になる。

曲率の影響

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

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

NOTE

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

まとめ：球面という設計の選択

ノルムを諦めるという一つの決断が、空間そのものを平面から球面へ作り替えた。プラネタリウムという比喩も、vMF分布という道具も、その決断の帰結にすぎない。

概念	定義	本回での役割
球面制約	すべての埋め込みを単位球面 ${\mathbb{S}}^{d-1}$ 上に置く	ノルムの混乱を根本から断ち切る設計
nGPT	重みも埋め込みも完全に正規化したTransformer設計	球面制約の実用例。学習効率の改善も報告
プラネタリウム・メタファー	球面をドーム、埋め込みを星として捉える	角度が意味を担うという直感の可視化
vMF分布	球面上の方向分布。集中度 $\kappa$ が分布の鋭さを制御	Softmaxとの形式的対応。第4回と第7回に接続
球面設計の限界	ノルムに意味がある場合、回帰タスク等では適さない	「球面が常に良い」という過剰一般化を防ぐ

球面制約は「長さと向きの混乱」を解消する強力な設計だが、あらゆる場面に効く万能薬ではない。ノルムが意味を持つタスクや、階層構造を表現したい場面では別の幾何学（第12回：双曲空間）が有効になる。どの空間を選ぶか、という問いは、本講義を通じて繰り返し問われる。

TIP

読了後に、理解度チェックで要点を確認できる。 このページで登場した用語は、キーワード集でも確認できる。

次回予告

第4回「分類を確率で読む」では、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., 2025）を参照されたい。

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. (2025). nGPT: Normalized Transformer with Representation Learning on the Hypersphere. ICLR 2025. arXiv: 2410.01131 (2024)
  • nGPTの原論文。Transformerのすべての表現を単位球面上に制約する設計。arXiv 2024、ICLR 2025に採択。
• Wang, F., Xiang, X., Cheng, J., & Yuille, A. L. (2017). NormFace: L2 Hypersphere Embedding for Face Verification. ACM MM 2017. arXiv: 1704.06369
  • 顔認証における球面埋め込みの先駆的研究。

von Mises-Fisher分布

• Mardia, K. V., & Jupp, P. E. (2000). Directional Statistics. Wiley. DOI: 10.1002/9780470316979
  • 方向統計学の標準的教科書。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. https://jmlr.org/papers/v6/banerjee05a.html
  • vMF分布を用いたクラスタリング。機械学習への応用。

リーマン最適化

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