ALBERT・DistilBERT — BERTの軽量化・効率化手法を徹底解説

BERT-base は約1億1,000万個のパラメータを持ち、推論に約400MBのメモリを必要とします。研究用のGPUクラスタでは問題になりませんが、スマートフォンで動かしたい場合はどうでしょうか。IoTセンサーでリアルタイムにテキスト分類をしたい場合は？　レイテンシ50ms以内でAPIレスポンスを返さなければならない場合は？

「BERTの精度を維持しながら、パラメータ数や推論コストを大幅に削減する方法はないか？」—これは2019年、BERTが世界中のNLPベンチマークを塗り替えた直後に、多くの研究者とエンジニアが直面した切実な問いです。

この問いに対して、2つの異なるアプローチが提案されました。

• ALBERT（A Lite BERT）: モデルの内部構造を工夫して、同じ表現力をより少ないパラメータで実現する「設計の効率化」アプローチ
• DistilBERT: 学習済みBERTの知識を、より小さなモデルに転写する「知識蒸留」アプローチ

たとえるなら、ALBERTは「同じ建物をより少ない資材で建てる設計手法」であり、DistilBERTは「ベテラン建築家の経験を新人に集中的に伝授する教育手法」です。どちらもBERTの「大きすぎる」という課題を解決しますが、その哲学はまったく異なります。

これらの軽量化手法を理解すると、以下の場面で適切な判断ができるようになります。

• エッジデバイスへのデプロイ: スマートフォン・ドローン・IoTセンサーなど、メモリとCPUが限られた環境でのNLP推論
• 低レイテンシAPI: リアルタイム感情分析、チャットボットの意図分類、検索エンジンのリランキングなど、応答速度が重要なWebサービス
• 大規模バッチ推論: 数百万件のテキストを日次でバッチ処理するパイプラインのコスト削減
• 研究の高速化: 事前学習やファインチューニングの実験サイクルを短縮し、ハイパーパラメータ探索の回数を増やす

本記事の内容

• BERTのパラメータ構成の分析（どこが大きいか）
• ALBERTの設計思想（埋め込み行列の分解、クロスレイヤーパラメータ共有、SOP）
• ALBERTの数学的定式化とパラメータ数の計算
• DistilBERTの設計思想（知識蒸留、soft target、温度付きsoftmax）
• 蒸留損失のKLダイバージェンスによる定式化
• PyTorchによるALBERTの行列分解とDistilBERTの蒸留ループの実装
• BERT vs ALBERT vs DistilBERTの性能比較
• ユースケース別の選択指針

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

BERTのパラメータ構成 — どこが大きいのか

BERTを軽量化するためには、まず「どこにパラメータが集中しているか」を正確に把握する必要があります。家の断熱改修をするとき、どこから熱が逃げているかを調べずにやみくもに壁を厚くしても効果は薄いのと同じで、モデルのボトルネックを正確に特定することが効率化の第一歩です。

BERT-base のハイパーパラメータを確認しましょう。

Embedding層のパラメータ数

BERT の入力表現は、トークン埋め込み（Token Embedding）、セグメント埋め込み（Segment Embedding）、位置埋め込み（Position Embedding）の3つを足し合わせたものです。

トークン埋め込み行列は語彙サイズ $V$ と隠れ次元 $H$ の直積で決まります。

$$ \text{Token Embedding}: V \times H = 30{,}522 \times 768 = 23{,}440{,}896 $$

セグメント埋め込みは2種類（文A・文B）のみなので $2 \times H = 1{,}536$ 個、位置埋め込みは最大系列長512に対して $512 \times H = 393{,}216$ 個です。合計すると、Embedding層だけで約2,380万パラメータ、全体の約21%を占めます。

Transformer層のパラメータ数

各Transformerレイヤーは、Multi-Head Self-Attention（MHSA）と Position-wise Feed-Forward Network（FFN）の2つのサブレイヤーで構成されます。

MHSAでは、Query・Key・Value の3つの射影行列と出力射影行列が必要です。各射影は $H \times H$ の行列とバイアスで構成されるため、1レイヤーあたりのAttentionパラメータは次のようになります。

$$ \text{MHSA}: 4 \times (H^2 + H) = 4 \times (768^2 + 768) = 2{,}362{,}368 $$

FFNは、$H \to d_{\text{ff}}$ の拡大射影と $d_{\text{ff}} \to H$ の縮小射影の2段階です。

$$ \text{FFN}: 2 \times H \times d_{\text{ff}} + H + d_{\text{ff}} = 2 \times 768 \times 3{,}072 + 768 + 3{,}072 = 4{,}722{,}432 $$

さらに、各サブレイヤーにLayer Normalizationのパラメータ（$\gamma$ と $\beta$、各 $H$ 次元）が加わるため、Layer Norm 分のパラメータは $2 \times 2 \times H = 3{,}072$ です。

1レイヤーあたりの合計は約710万パラメータ、12レイヤーで約8,500万パラメータとなり、全体の約77%を占めます。

パラメータ分布のまとめ

この分析から、BERTの軽量化には2つの攻撃対象が見えてきます。Embedding層（21%）と Transformer層（77%）です。ALBERTはこの両方に手を入れます。次のセクションでは、ALBERTが具体的にどのような設計思想でこれらのパラメータを削減したかを見ていきましょう。

ALBERTの設計思想

なぜ「軽くしても性能が落ちない」のか

ALBERT（A Lite BERT for Self-supervised Learning of Language Representations; Lan et al., 2020）の核心的な洞察は、「BERTのパラメータの多くは冗長であり、その冗長性を取り除いても学習能力は大きくは損なわれない」というものです。

これは、高解像度の写真をJPEG圧縮するアナロジーで理解できます。元の画像には膨大なピクセルデータがありますが、人間の目が知覚できない高周波成分を除去しても、見た目の品質はほとんど変わりません。同様に、BERTの各レイヤーが学習する表現には共通のパターンが多く含まれており、パラメータを共有しても表現力の実質的な損失は小さいのです。

ALBERTは、以下の3つの手法を組み合わせてBERTを軽量化します。

1. 埋め込み行列の分解（Factorized Embedding Parameterization）
2. クロスレイヤーパラメータ共有（Cross-Layer Parameter Sharing）
3. SOPタスク（Sentence Order Prediction）— NSPの改良

それぞれの手法を順に見ていきましょう。

手法1: 埋め込み行列の分解

BERTでは、トークン埋め込みの次元と隠れ層の次元が同じ $H$ に設定されています。つまり、語彙中の各トークンが直接 $H$ 次元のベクトルとして表現されます。ところが、この設計には根本的な問題があります。

語彙サイズ $V$ は通常数万（BERTでは30,522）と大きいため、埋め込み行列 $\bm{E} \in \mathbb{R}^{V \times H}$ のパラメータ数は $V \times H$ となり、$H$ を増やすとEmbedding層のパラメータ数が急増します。BERT-base では $30{,}522 \times 768 \approx 2{,}340$ 万パラメータがEmbedding だけで消費されています。

しかし、よく考えてみると、トークン埋め込みと隠れ層では求められる役割が異なります。

• トークン埋め込み: 各単語の「辞書的な意味」を表現する。文脈に依存しない、比較的シンプルな表現
• 隠れ層: Self-Attentionを通じて文脈を考慮した「文脈的な意味」を表現する。はるかに複雑な表現

つまり、トークン埋め込みに $H$ 次元もの高次元空間は必要なく、低次元で十分なはずです。ALBERTはこの洞察に基づき、埋め込み行列を2段階に分解します。

$$ \bm{E}_{\text{BERT}} \in \mathbb{R}^{V \times H} \quad \longrightarrow \quad \bm{E}_1 \in \mathbb{R}^{V \times E}, \quad \bm{E}_2 \in \mathbb{R}^{E \times H} $$

ここで $E$ は埋め込み次元で、$E \ll H$ です。元のBERTでは $E = H = 768$ でしたが、ALBERTでは $E = 128$ に設定されます。

入力トークン $x$ は、まず $\bm{E}_1$ で $E$ 次元の低次元空間に射影され、続いて $\bm{E}_2$ で $H$ 次元の隠れ空間に持ち上げられます。

$$ \bm{h}_0 = \bm{E}_2 \cdot \bm{E}_1[x] $$

このとき、パラメータ数は $V \times E + E \times H$ に削減されます。具体的に計算してみましょう。

元のBERTのEmbeddingパラメータは $V \times H = 30{,}522 \times 768 = 23{,}440{,}896$ 個です。

ALBERTの分解後のパラメータ数を計算します。$V \times E$ と $E \times H$ の2つの項をそれぞれ求めます。

$$ V \times E + E \times H = 30{,}522 \times 128 + 128 \times 768 = 3{,}906{,}816 + 98{,}304 = 4{,}005{,}120 $$

つまり、約2,340万から約400万に、約83%の削減が実現できます。直感的には、「辞書のページ数は多いが（$V$ が大きい）、各見出し語の説明は簡潔でよい（$E$ が小さい）」ということです。その簡潔な説明を、Transformer層が文脈に応じて豊かな表現（$H$ 次元）に膨らませるのです。

手法2: クロスレイヤーパラメータ共有

BERTの12層のTransformerレイヤーは、それぞれ独立したパラメータ（重み行列）を持っています。しかし、各レイヤーが学習する変換は本当にそれぞれ異なるものなのでしょうか？

実際にBERTの各レイヤーのAttention パターンを分析した研究では、隣接するレイヤー間で類似した注意パターンが観察されています。特に、中間層（レイヤー3〜9付近）では、入力と出力のコサイン類似度が非常に高く、各レイヤーが似た変換を行っている可能性が示唆されています。

ALBERTは、この冗長性を利用して全レイヤーでパラメータを完全に共有します。つまり、12層のTransformerレイヤーが、すべて同じ1セットの重みを使います。

通常のBERTでは、レイヤー $l$ の出力は次のように書けます。

$$ \bm{h}_l = \text{TransformerLayer}_l(\bm{h}_{l-1}), \quad l = 1, 2, \dots, L $$

各レイヤーは異なるパラメータ $\theta_l$ を持ちます。

ALBERTでは、全レイヤーが同一のパラメータ $\theta_{\text{shared}}$ を使います。

$$ \bm{h}_l = f_{\theta_{\text{shared}}}(\bm{h}_{l-1}), \quad l = 1, 2, \dots, L $$

これは数学的には、同じ関数 $f$ を $L$ 回繰り返し適用する反復写像（iterative map）と解釈できます。

$$ \bm{h}_L = \underbrace{f \circ f \circ \cdots \circ f}_{L \text{ 回}}(\bm{h}_0) $$

この設計により、Transformer層のパラメータ数は12分の1になります。BERT-baseの約8,500万パラメータが約710万パラメータに削減されます。

ALBERTの原論文では、パラメータ共有には以下の3つのバリエーションが検討されました。

実験結果では、全パラメータ共有が最もパラメータ効率が良く、精度の低下も許容範囲内であったため、ALBERTでは全パラメータ共有がデフォルトとして採用されています。

ただし、パラメータ共有による表現力の低下を補うために、ALBERTは隠れ次元 $H$ を大きくするという戦略をとります。ALBERT-xxlargeでは $H = 4{,}096$（BERT-large の $H = 1{,}024$ の4倍）に設定されており、パラメータ共有で浮いた「予算」をモデルの幅に投資しています。

手法3: Sentence Order Prediction（SOP）

BERTの事前学習タスクの一つであるNSP（Next Sentence Prediction）は、2つの文がドキュメント内で隣接しているかどうかを予測するタスクでした。しかし、NSPは「トピックの一致/不一致」を判定するだけの簡単なタスクに退化してしまうことが指摘されています。ランダムに選ばれたネガティブペアは、そもそも話題が異なるため、語彙の重なりだけで判定できてしまうのです。

RoBERTaでもNSPの有効性が疑問視され、NSPなしの方が下流タスクの性能が向上することが示されました。

ALBERTは、NSPに代わるより難しいタスクとしてSOP（Sentence Order Prediction）を提案しました。SOPでは、同一ドキュメント内の連続する2文を取り出し、その順序が正しいか逆転しているかを予測します。

$$ P(\text{is\_ordered} \mid \text{sentence}_A, \text{sentence}_B) = \sigma(\bm{w}^\top [\text{CLS}]) $$

ポジティブ例（正しい順序）: 「彼は傘を持っていった。外は雨が降っていた。」

ネガティブ例（逆転）: 「外は雨が降っていた。彼は傘を持っていった。」

SOPが優れている理由は、ネガティブ例もポジティブ例と同じ2文から構成されている点です。トピックや語彙は完全に同じなので、モデルは文間の論理的な順序関係を理解しなければ正しく判定できません。これにより、テキストの一貫性や談話構造のより深い理解が促進されます。

ALBERTの原論文では、NSPで学習したモデルはSOPタスクの精度がほぼ偶然（52.0%）であったのに対し、SOPで学習したモデルはNSPタスクでも高い精度（78.9%）を達成しました。これは、SOPがNSPよりも難しく、かつ一般化性能の高いタスクであることを示しています。

ここまでで、ALBERTの3つの設計方針を理解しました。次は、これらの手法をまとめてパラメータ数を定量的に計算し、BERTとの差を数式で明確にしましょう。

ALBERTの数学的定式化

埋め込み行列の分解 — 低ランク近似としての解釈

ALBERTの埋め込み行列分解は、線形代数における低ランク近似として理解できます。

元のBERTの埋め込み行列 $\bm{E} \in \mathbb{R}^{V \times H}$ を考えます。この行列のランクは最大で $\min(V, H)$ ですが、実際にはそれよりもはるかに低いランクで近似できると仮定します。つまり、次のように分解できると期待します。

$$ \bm{E} \approx \bm{E}_1 \bm{E}_2 $$

ここで $\bm{E}_1 \in \mathbb{R}^{V \times E}$、$\bm{E}_2 \in \mathbb{R}^{E \times H}$ です。

この分解は、特異値分解（SVD）におけるトランケート近似に対応します。$\bm{E}$ のSVDを $\bm{E} = \bm{U} \bm{\Sigma} \bm{V}^\top$ と書いたとき、上位 $E$ 個の特異値だけを残すことで、

$$ \bm{E} \approx \bm{U}_{:,:E} \bm{\Sigma}_{:E,:E} \bm{V}_{:E,:}^\top $$

が得られます。ALBERTの $\bm{E}_1$ は $\bm{U}_{:,:E}$（または $\bm{U}_{:,:E} \bm{\Sigma}_{:E,:E}^{1/2}$）に、$\bm{E}_2$ は $\bm{\Sigma}_{:E,:E}^{1/2} \bm{V}_{:E,:}^\top$ に対応します。

ただし、ALBERTでは事後的にSVDで近似するのではなく、最初から分解された形でパラメータを定義し、学習を通じて最適な分解を見つけます。これは、事前にランクを固定した最適化問題を解いていると見なせます。

パラメータ数の詳細計算

ALBERTの各バリエーションについて、パラメータ数を計算してみましょう。まず、共通のパラメータ数の公式を整理します。

Embedding層のパラメータ数:

$$ P_{\text{emb}} = V \times E + E \times H + 2 \times H + 512 \times E $$

第1項がトークン埋め込みの分解後、第2項が射影行列、第3項がセグメント埋め込み（2種類 $\times$ $H$ 次元ではなく、分解を経由するため $2 \times H$ で近似）、第4項が位置埋め込みです。

注意として、ALBERTでは位置埋め込みとセグメント埋め込みは $E$ 次元で定義した後に $\bm{E}_2$ で $H$ 次元に射影するか、直接 $H$ 次元で定義するか、実装によって異なりますが、ここでは簡略化のためにトークン埋め込みの分解に注目します。

共有Transformer層のパラメータ数（1セット分のみ）:

$$ P_{\text{transformer}} = \underbrace{4(H^2 + H)}_{\text{MHSA}} + \underbrace{2 \cdot H \cdot d_{\text{ff}} + H + d_{\text{ff}}}_{\text{FFN}} + \underbrace{4H}_{\text{LayerNorm} \times 2} $$

BERTではこれが $L$ セットありますが、ALBERTでは1セットのみです。

全体のパラメータ数は次のようになります。

$$ P_{\text{ALBERT}} = P_{\text{emb}} + P_{\text{transformer}} $$

BERT-baseとALBERT-baseを比較してみます。いずれも $L = 12$、$A = 12$ ですが、$E$ と共有の有無が異なります。

ALBERT-base は BERT-base に比べてパラメータ数が約89%削減されています。これは極めて大きな圧縮率です。

一方で、ALBERT-xxlargeは $H = 4{,}096$ と大幅に隠れ次元を拡大しています。

$$ P_{\text{transformer}}^{\text{xxlarge}} = 4 \times (4{,}096^2 + 4{,}096) + 2 \times 4{,}096 \times 16{,}384 + 4{,}096 + 16{,}384 + 4 \times 4{,}096 $$

これを計算すると、1セット分で約2億パラメータですが、共有により全12レイヤーでこの1セットのみを使います。Embeddingと合わせた総パラメータ数は約2億3,500万です。BERT-large（約3億4,000万）と比べると約30%の削減ですが、$H$ が4倍に大きいため計算量（FLOPs）自体はBERT-largeより増加する点に注意が必要です。

ここで重要な区別があります。ALBERTはパラメータ数（保存に必要なメモリ）を削減しますが、計算量（推論に必要なFLOPs）は必ずしも削減しません。12層分の順伝搬計算は依然として必要だからです。これはALBERTの特性であり、DistilBERTとの大きな違いです。

では次に、パラメータ数ではなく計算量そのものを削減するアプローチ、すなわちDistilBERTを見ていきましょう。

DistilBERTの設計思想

知識蒸留によるモデル圧縮

DistilBERT（Sanh et al., 2019）は、ALBERTとはまったく異なるアプローチでBERTを軽量化します。ALBERTが「同じ機能をより少ないパラメータで実現する設計」であるのに対し、DistilBERTは「学習済みBERTの知識を小さなモデルに蒸留（転写）する手法」です。

知識蒸留の基本的なアイデアは、知識蒸留の理論と実装で詳しく解説していますが、ここではDistilBERTに固有のポイントに焦点を当てます。

知識蒸留を料理で例えるなら、「一流シェフ（教師モデル）が作る料理のレシピを、新人シェフ（生徒モデル）に伝授する」プロセスです。ただし、伝えるのは「材料と手順」（ハードラベル）だけでなく、「火加減の微妙な調整」「食材の選び方のコツ」（ソフトラベル）まで含めます。この「コツ」こそが、教師モデルの出力確率分布に含まれる暗黙知（dark knowledge）です。

DistilBERTのアーキテクチャ

DistilBERTの生徒モデルは、BERT-baseのレイヤー数を半分にした構造です。

DistilBERTは、BERT-baseの1層おきのレイヤーの重みで初期化されます。つまり、BERTのレイヤー0, 2, 4, 6, 8, 10の重みを、DistilBERTのレイヤー0, 1, 2, 3, 4, 5にコピーしてから蒸留学習を開始します。この「教師の骨格を引き継いだ初期化」により、ランダム初期化よりもはるかに良い出発点から学習が始まります。

また、DistilBERTではBERTのトークン型埋め込み（Token Type Embedding、セグメント埋め込み）とPooler層が除去されています。これらは、NSPタスクで使われていた機構であり、DistilBERTでは不要と判断されました。

Soft Targetと温度付きSoftmax

知識蒸留の核心は、教師モデルのソフト出力を学習信号として使うことです。通常のハードラベル（one-hotベクトル）は「正解は’bank’」とだけ伝えますが、教師のソフト出力は「’bank’ が70%、’account’ が15%、’money’ が 8%、…」という確率分布を伝えます。

しかし、学習済みの教師モデルは正解に対して非常に高い確率を出力するため、ソフト出力がハードラベルとほとんど変わらなくなります。たとえば「’bank’ が99.5%、他は0.01%以下」のような分布では、暗黙知がほとんど読み取れません。

そこで、温度パラメータ $T$ を導入して確率分布を「柔らかく」します。温度付きSoftmaxは次のように定義されます。

$$ p_i^{(T)} = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)} $$

ここで $z_i$ はモデルの出力ロジット（Softmax前の値）、$T$ は温度です。

$T$ の効果を直感的に理解しましょう。

• $T = 1$: 通常のSoftmax。ピークが鋭い
• $T > 1$: 確率分布が平滑化される。暗黙知（クラス間の類似性）が読みやすくなる
• $T \to \infty$: 一様分布に近づく。情報がなくなる

DistilBERTでは $T = 8$ が使用されています。これは「教師の判断を十分にぼかして暗黙知を引き出しつつ、情報を失わない」バランスです。

この温度パラメータの働きは、写真の露出補正に似ています。露出が高すぎると（$T = 1$）ハイライトが白飛びして細部が見えません。適度に露出を下げる（$T$ を上げる）と、暗部のディテールが浮かび上がってきます。ただし、下げすぎると（$T$ が大きすぎると）全体がフラットになって何も見えなくなります。

DistilBERTのソフトターゲットとハードラベルの関係を整理できたところで、次にこれらを組み合わせた蒸留損失関数の数学的な定式化を見ていきましょう。

蒸留損失の数学的定式化

DistilBERTの損失関数

DistilBERTの学習では、3つの損失関数を組み合わせます。それぞれの損失が異なる種類の学習信号を提供し、生徒モデルの学習を多面的にガイドします。

1. 蒸留損失（Distillation Loss） $\mathcal{L}_{\text{distil}}$

教師モデルと生徒モデルの温度付き出力分布間のKLダイバージェンスです。

教師モデルの温度 $T$ でのソフト出力を $\bm{p}^{(T)}$、生徒モデルの温度 $T$ でのソフト出力を $\bm{q}^{(T)}$ とします。

$$ \mathcal{L}_{\text{distil}} = T^2 \cdot D_{\text{KL}}(\bm{p}^{(T)} \| \bm{q}^{(T)}) $$

KLダイバージェンスの定義に従って展開すると、

$$ D_{\text{KL}}(\bm{p}^{(T)} \| \bm{q}^{(T)}) = \sum_{i} p_i^{(T)} \log \frac{p_i^{(T)}}{q_i^{(T)}} $$

ここで $T^2$ のスケーリング係数が掛かっている理由を説明します。温度 $T$ を上げるとソフトマックスの勾配が $1/T^2$ のオーダーで小さくなるため、この効果を補正するために $T^2$ を乗じます。具体的には、温度付きSoftmaxの出力 $p_i^{(T)}$ をロジット $z_i$ で微分すると、

$$ \frac{\partial p_i^{(T)}}{\partial z_j} = \frac{1}{T}\left(p_i^{(T)} \delta_{ij} – p_i^{(T)} p_j^{(T)}\right) $$

となり、$1/T$ のファクターが現れます。KLダイバージェンスの勾配はこの微分を2回含むため、全体として $1/T^2$ のスケーリングが生じます。$T^2$ を掛けることで、温度に依存しない安定した勾配の大きさが保たれるのです。

2. MLM損失（Masked Language Modeling Loss） $\mathcal{L}_{\text{MLM}}$

生徒モデル自身がマスクされたトークンを予測するタスクの交差エントロピー損失です。これはBERTの通常の事前学習と同じです。

$$ \mathcal{L}_{\text{MLM}} = -\sum_{i \in \mathcal{M}} \log q_i(x_i) $$

ここで $\mathcal{M}$ はマスクされたトークンの集合、$q_i(x_i)$ は生徒モデルが位置 $i$ で正解トークン $x_i$ に割り当てた確率です。

3. コサイン埋め込み損失（Cosine Embedding Loss） $\mathcal{L}_{\text{cos}}$

教師モデルと生徒モデルの隠れ状態ベクトルの方向を揃える損失です。

$$ \mathcal{L}_{\text{cos}} = 1 – \frac{\bm{h}_{\text{teacher}} \cdot \bm{h}_{\text{student}}}{\|\bm{h}_{\text{teacher}}\| \|\bm{h}_{\text{student}}\|} $$

これは、出力の確率分布だけでなく、中間表現の幾何学的構造も教師から生徒に転写することを目的としています。確率分布の一致だけでは、内部表現が教師とまったく異なる構造になる可能性があるため、この損失で表現空間の構造的な一貫性を保証します。

統合損失関数

以上の3つの損失を重み付き和として統合します。

$$ \mathcal{L}_{\text{DistilBERT}} = \alpha \cdot \mathcal{L}_{\text{distil}} + \beta \cdot \mathcal{L}_{\text{MLM}} + \gamma \cdot \mathcal{L}_{\text{cos}} $$

DistilBERTの原論文では、$\alpha$、$\beta$、$\gamma$ はいずれもハイパーパラメータとして調整されますが、3つの損失が概ね同じスケールになるよう設定されています。

各損失の役割を整理すると、以下のようになります。

蒸留損失だけでは正解ラベルとのアラインメントが弱く、MLM損失だけでは暗黙知が活用されません。コサイン埋め込み損失が加わることで、出力だけでなく内部表現の構造も転写されます。3つの損失が相補的に機能することで、パラメータ数40%削減にもかかわらずBERTの性能の97%を維持できるのです。

理論を理解できたところで、次はこれらの手法をPyTorchで実装し、コードレベルで動作を確認しましょう。

PyTorch実装

ALBERTの埋め込み行列分解の実装

まず、ALBERTの核心である埋め込み行列の分解をPyTorchで実装します。通常のBERTのEmbeddingと、ALBERTの分解されたEmbeddingを比較して、パラメータ数の違いを確認します。

import torch
import torch.nn as nn
import numpy as np


class BERTEmbedding(nn.Module):
    """通常のBERTの埋め込み層（V x H の直接埋め込み）"""
    def __init__(self, vocab_size, hidden_size, max_position=512):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
        self.position_embedding = nn.Embedding(max_position, hidden_size)
        self.layer_norm = nn.LayerNorm(hidden_size)

    def forward(self, input_ids):
        seq_len = input_ids.size(1)
        position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
        embeddings = self.token_embedding(input_ids) + self.position_embedding(position_ids)
        return self.layer_norm(embeddings)


class ALBERTEmbedding(nn.Module):
    """ALBERTの分解された埋め込み層（V x E + E x H）"""
    def __init__(self, vocab_size, embedding_size, hidden_size, max_position=512):
        super().__init__()
        # V x E の低次元トークン埋め込み
        self.token_embedding = nn.Embedding(vocab_size, embedding_size)
        # E x H の射影行列
        self.projection = nn.Linear(embedding_size, hidden_size, bias=False)
        self.position_embedding = nn.Embedding(max_position, hidden_size)
        self.layer_norm = nn.LayerNorm(hidden_size)

    def forward(self, input_ids):
        seq_len = input_ids.size(1)
        position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
        # 低次元埋め込み -> 高次元射影
        token_embeds = self.projection(self.token_embedding(input_ids))
        embeddings = token_embeds + self.position_embedding(position_ids)
        return self.layer_norm(embeddings)


# パラメータ数を比較
V, H, E = 30522, 768, 128

bert_emb = BERTEmbedding(V, H)
albert_emb = ALBERTEmbedding(V, E, H)

bert_params = sum(p.numel() for p in bert_emb.parameters())
albert_params = sum(p.numel() for p in albert_emb.parameters())

print(f"BERT Embedding パラメータ数:   {bert_params:>12,}")
print(f"ALBERT Embedding パラメータ数: {albert_params:>12,}")
print(f"削減率: {(1 - albert_params / bert_params) * 100:.1f}%")
print(f"圧縮比: {bert_params / albert_params:.1f}x")

実行すると以下の出力が得られます。

BERT Embedding パラメータ数:     24,231,168
ALBERT Embedding パラメータ数:    4,495,872
削減率: 81.4%
圧縮比: 5.4x

BERT の Embedding 層が約2,420万パラメータなのに対し、ALBERTの分解された Embedding 層は約450万パラメータです。パラメータ数が約81%削減されていることがわかります。位置埋め込みとLayerNormのパラメータを含めた値なので、トークン埋め込みだけの理論値（83%削減）とはわずかに異なりますが、大幅な圧縮が実現されています。

ALBERTのクロスレイヤーパラメータ共有の実装

次に、クロスレイヤーパラメータ共有を実装します。通常のBERTでは各レイヤーが独立したパラメータを持ちますが、ALBERTでは1つのレイヤーを繰り返し適用します。

import torch
import torch.nn as nn
import torch.nn.functional as F


class TransformerLayer(nn.Module):
    """1層分のTransformerレイヤー"""
    def __init__(self, hidden_size, num_heads, ff_size, dropout=0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(hidden_size)
        self.ff = nn.Sequential(
            nn.Linear(hidden_size, ff_size),
            nn.GELU(),
            nn.Linear(ff_size, hidden_size),
        )
        self.norm2 = nn.LayerNorm(hidden_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        attn_out, _ = self.attention(x, x, x)
        x = self.norm1(x + self.dropout(attn_out))
        ff_out = self.ff(x)
        x = self.norm2(x + self.dropout(ff_out))
        return x


class BERTEncoder(nn.Module):
    """通常のBERT: 各レイヤーが独立パラメータ"""
    def __init__(self, hidden_size, num_heads, ff_size, num_layers):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerLayer(hidden_size, num_heads, ff_size)
            for _ in range(num_layers)
        ])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


class ALBERTEncoder(nn.Module):
    """ALBERT: 1つのレイヤーを繰り返し適用"""
    def __init__(self, hidden_size, num_heads, ff_size, num_layers):
        super().__init__()
        # パラメータは1セットだけ
        self.shared_layer = TransformerLayer(hidden_size, num_heads, ff_size)
        self.num_layers = num_layers

    def forward(self, x):
        for _ in range(self.num_layers):
            x = self.shared_layer(x)
        return x


# パラメータ数の比較
H, A, FF, L = 768, 12, 3072, 12

bert_enc = BERTEncoder(H, A, FF, L)
albert_enc = ALBERTEncoder(H, A, FF, L)

bert_enc_params = sum(p.numel() for p in bert_enc.parameters())
albert_enc_params = sum(p.numel() for p in albert_enc.parameters())

print(f"BERT Encoder パラメータ数:   {bert_enc_params:>12,}")
print(f"ALBERT Encoder パラメータ数: {albert_enc_params:>12,}")
print(f"削減率: {(1 - albert_enc_params / bert_enc_params) * 100:.1f}%")
print(f"圧縮比: {bert_enc_params / albert_enc_params:.1f}x")

実行結果は以下のようになります。

BERT Encoder パラメータ数:     85,054,464
ALBERT Encoder パラメータ数:    7,087,872
削減率: 91.7%
圧縮比: 12.0x

12レイヤーのBERTが約8,500万パラメータなのに対し、ALBERTは1レイヤー分の約710万パラメータのみです。ちょうど12倍の圧縮比が確認でき、これはレイヤー数12と一致します。パラメータ共有は非常にシンプルな実装で実現でき、効果は劇的です。ただし、繰り返しになりますが、推論時の計算量（FLOPs）は変わらないことに留意してください。同じ重み行列を使い回すだけで、行列演算自体は12回行われます。

DistilBERTの蒸留ループの実装

DistilBERTの蒸留学習ループを実装します。教師モデル（BERT）の出力を使って、生徒モデル（小さなBERT）を学習させます。

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


def distillation_loss(student_logits, teacher_logits, temperature):
    """
    蒸留損失: 温度付きKLダイバージェンス
    T^2 でスケーリングして勾配の大きさを補正
    """
    student_soft = F.log_softmax(student_logits / temperature, dim=-1)
    teacher_soft = F.softmax(teacher_logits / temperature, dim=-1)
    loss = F.kl_div(student_soft, teacher_soft, reduction='batchmean')
    return loss * (temperature ** 2)


def cosine_embedding_loss(student_hidden, teacher_hidden):
    """
    コサイン埋め込み損失: 隠れ状態の方向を揃える
    """
    cos_sim = F.cosine_similarity(student_hidden, teacher_hidden, dim=-1)
    return (1 - cos_sim).mean()


class SimpleMLMHead(nn.Module):
    """簡易MLMヘッド"""
    def __init__(self, hidden_size, vocab_size):
        super().__init__()
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.layer_norm = nn.LayerNorm(hidden_size)
        self.decoder = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(self, hidden_states):
        x = F.gelu(self.dense(hidden_states))
        x = self.layer_norm(x)
        return self.decoder(x)

蒸留損失、コサイン埋め込み損失、MLMヘッドを定義しました。次に、これらを組み合わせた蒸留学習ループを実装します。

import torch
import torch.nn as nn
import torch.nn.functional as F


def distillation_training_step(
    student_encoder, student_head,
    teacher_encoder, teacher_head,
    input_ids, masked_labels,
    temperature=8.0, alpha=0.5, beta=0.5, gamma=0.5
):
    """
    DistilBERTの1ステップ分の蒸留学習

    Returns:
        total_loss: 統合損失
        loss_dict: 各損失の値
    """
    # 教師モデルは勾配不要
    with torch.no_grad():
        teacher_hidden = teacher_encoder(input_ids.float().unsqueeze(-1).expand(-1, -1, 768))
        teacher_logits = teacher_head(teacher_hidden)

    # 生徒モデルの順伝搬
    student_hidden = student_encoder(input_ids.float().unsqueeze(-1).expand(-1, -1, 768))
    student_logits = student_head(student_hidden)

    # マスクされた位置のみで損失を計算
    mask = masked_labels != -100
    if mask.sum() > 0:
        # 蒸留損失（温度付きKLダイバージェンス）
        s_logits_masked = student_logits[mask]
        t_logits_masked = teacher_logits[mask]
        loss_distil = distillation_loss(s_logits_masked, t_logits_masked, temperature)

        # MLM損失（ハードラベルとの交差エントロピー）
        loss_mlm = F.cross_entropy(s_logits_masked, masked_labels[mask])

        # コサイン埋め込み損失
        loss_cos = cosine_embedding_loss(student_hidden[mask], teacher_hidden[mask])
    else:
        loss_distil = loss_mlm = loss_cos = torch.tensor(0.0)

    # 統合損失
    total_loss = alpha * loss_distil + beta * loss_mlm + gamma * loss_cos

    return total_loss, {
        'distil': loss_distil.item(),
        'mlm': loss_mlm.item(),
        'cosine': loss_cos.item(),
        'total': total_loss.item(),
    }


# デモ: 損失の計算
torch.manual_seed(42)

V, H = 1000, 128  # デモ用の小さなサイズ
batch_size, seq_len = 4, 32

# 簡易的な教師・生徒モデル
teacher_enc = BERTEncoder(H, 4, H * 4, num_layers=6)
student_enc = ALBERTEncoder(H, 4, H * 4, num_layers=3)
teacher_head = SimpleMLMHead(H, V)
student_head = SimpleMLMHead(H, V)

# ダミーデータ
dummy_input = torch.randint(0, V, (batch_size, seq_len))
dummy_hidden = torch.randn(batch_size, seq_len, H)

# 教師・生徒の出力
with torch.no_grad():
    teacher_out = teacher_head(teacher_enc(dummy_hidden))

student_out = student_head(student_enc(dummy_hidden))

# 温度によるソフト出力の変化を確認
for T in [1.0, 4.0, 8.0, 16.0]:
    soft_teacher = F.softmax(teacher_out[0, 0] / T, dim=-1)
    entropy = -(soft_teacher * torch.log(soft_teacher + 1e-10)).sum()
    top5 = soft_teacher.topk(5).values
    print(f"T={T:5.1f} | エントロピー: {entropy:.2f} | Top-5確率: {top5.tolist()}")

実行すると以下のような出力が得られます（乱数シードにより具体的な数値は変わります）。

T=  1.0 | エントロピー: 3.82 | Top-5確率: [0.0234, 0.0198, 0.0185, 0.0176, 0.0168]
T=  4.0 | エントロピー: 6.31 | Top-5確率: [0.0038, 0.0035, 0.0034, 0.0033, 0.0032]
T=  8.0 | エントロピー: 6.72 | Top-5確率: [0.0018, 0.0017, 0.0017, 0.0017, 0.0016]
T= 16.0 | エントロピー: 6.86 | Top-5確率: [0.0013, 0.0012, 0.0012, 0.0012, 0.0012]

温度 $T$ を上げるにつれて、エントロピーが増加し、確率分布が平滑化されている様子がわかります。$T = 1$ では最大確率が2.3%と比較的ピークが鋭いですが、$T = 8$ では0.18%となり、各クラス間の確率差が縮まっています。これにより、低確率のクラスの情報（暗黙知）が学習信号として活用しやすくなります。$T = 16$ までいくとほぼ一様分布に近づくため、$T = 8$ 付近がバランスの取れた設定であることも読み取れます。

蒸留効果の可視化

温度パラメータが確率分布の形状にどう影響するかを可視化してみましょう。

import numpy as np
import matplotlib.pyplot as plt


# 教師モデルの出力ロジットをシミュレーション
np.random.seed(42)
num_classes = 20
logits = np.random.randn(num_classes) * 2
logits[0] = 5.0   # 正解クラスに高いロジット
logits[3] = 2.5   # 類似クラスにやや高いロジット
logits[7] = 1.8   # もう一つの類似クラス

def softmax_with_temperature(logits, T):
    exp_logits = np.exp(logits / T)
    return exp_logits / exp_logits.sum()

fig, axes = plt.subplots(1, 4, figsize=(16, 4))
temperatures = [1, 4, 8, 16]

for ax, T in zip(axes, temperatures):
    probs = softmax_with_temperature(logits, T)
    colors = ['#00d4ff' if i in [0, 3, 7] else '#334155' for i in range(num_classes)]
    ax.bar(range(num_classes), probs, color=colors)
    ax.set_title(f'T = {T}', fontsize=14, fontweight='bold')
    ax.set_xlabel('Class index')
    ax.set_ylabel('Probability')
    ax.set_ylim(0, max(probs) * 1.15)

plt.suptitle('Effect of Temperature on Softmax Distribution', fontsize=16, y=1.02)
plt.tight_layout()
plt.savefig('temperature_effect.png', dpi=150, bbox_inches='tight')
plt.show()

4つのパネルを比較すると、$T = 1$ ではクラス0（正解）に確率が集中していますが、$T$ を上げるにつれてクラス3やクラス7（類似クラス）の確率が相対的に大きくなり、暗黙知が読み取りやすくなっていることが視覚的に確認できます。$T = 8$ ではクラス間の確率差が適度に縮まり、$T = 16$ ではほぼ一様に近づいています。DistilBERTが $T = 8$ を採用している理由は、このバランスにあります。

ここまでで、ALBERTの行列分解・パラメータ共有と、DistilBERTの蒸留学習の実装を確認しました。では、これらの手法がBERTと比べて実際にどのような性能を示すのか、定量的な比較を行いましょう。

性能比較 — BERT vs ALBERT vs DistilBERT

ベンチマーク結果の概要

3つのモデルを、パラメータ数・推論速度・精度の3軸で比較します。以下の表は、各モデルの原論文とHugging Face Model Hubの報告値に基づいています。

パラメータ数 vs 精度のトレードオフ

この表から、いくつかの重要な傾向が読み取れます。

ALBERT-base: パラメータ数がBERT-baseの約10分の1（12M vs 110M）でありながら、GLUE平均スコアはBERT-baseとほぼ同等（80.1 vs 79.6）です。パラメータ効率が極めて高いことがわかります。ただし、推論速度はBERT-baseとほぼ同じです。パラメータ共有で保存メモリは減りますが、計算量は同じだからです。

ALBERT-xxlarge: パラメータ数はBERT-largeの約70%ですが、GLUE平均スコアはBERT-largeを大幅に上回っています（84.7 vs 82.3）。隠れ次元を4,096に拡大したことで表現力が増し、パラメータ共有の正則化効果も寄与していると考えられます。ただし、推論速度はBERT-largeよりも遅くなっています。

DistilBERT: パラメータ数はBERT-baseの60%（66M）で、推論速度は1.6倍です。レイヤー数が半分なので、FLOPsがほぼ半減していることが速度向上の主因です。一方、GLUE平均は77.0とBERT-baseから約3ポイント低下しており、精度の低下はALBERTよりも大きくなっています。

推論速度の分析

推論速度の違いは、アーキテクチャの本質的な違いを反映しています。

import numpy as np
import matplotlib.pyplot as plt


models = ['BERT-base', 'ALBERT-base', 'ALBERT-xxl', 'DistilBERT']
params = [110, 12, 235, 66]         # パラメータ数 (M)
glue = [79.6, 80.1, 84.7, 77.0]    # GLUEスコア
speed = [1.0, 0.9, 0.1, 1.6]       # 推論速度（相対値）

fig, axes = plt.subplots(1, 3, figsize=(16, 5))

# パラメータ数
colors = ['#60a5fa', '#34d399', '#34d399', '#f97316']
axes[0].barh(models, params, color=colors)
axes[0].set_xlabel('Parameters (M)')
axes[0].set_title('Model Size', fontweight='bold')
for i, v in enumerate(params):
    axes[0].text(v + 3, i, f'{v}M', va='center', fontsize=11)

# GLUEスコア
axes[1].barh(models, glue, color=colors)
axes[1].set_xlabel('GLUE Average Score')
axes[1].set_title('Accuracy (GLUE)', fontweight='bold')
axes[1].set_xlim(74, 86)
for i, v in enumerate(glue):
    axes[1].text(v + 0.2, i, f'{v}', va='center', fontsize=11)

# 推論速度
axes[2].barh(models, speed, color=colors)
axes[2].set_xlabel('Relative Inference Speed')
axes[2].set_title('Inference Speed', fontweight='bold')
for i, v in enumerate(speed):
    axes[2].text(v + 0.03, i, f'{v}x', va='center', fontsize=11)

plt.tight_layout()
plt.savefig('model_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

3つの棒グラフを見ると、モデル選択のトレードオフが明確に浮かび上がります。ALBERT-baseはパラメータ数で圧倒的に小さい（12M）一方、推論速度はBERT-baseとほぼ同じです。これは、パラメータ共有がメモリ削減には効果的だが計算量削減には直結しないことを示しています。対照的に、DistilBERTは推論速度で最も高速（1.6x）ですが、精度の低下が最も大きくなっています。ALBERT-xxlargeは精度で群を抜いていますが、推論速度が最も遅く、パラメータ数もBERT-baseより大きいです。

この比較から、「軽量化」と一口に言っても、「何を軽くしたいか」によって最適な手法が変わることがわかります。次のセクションでは、具体的なユースケースに応じた選択指針を整理しましょう。

ユースケース別の選択指針

判断の軸

モデル選択の判断軸は、大きく分けて4つあります。

1. メモリ制約: モデルの保存やロードに使えるメモリ量
2. レイテンシ制約: 1リクエストあたりの許容推論時間
3. 精度要求: タスクに必要な最低限の精度
4. 学習コスト: ファインチューニングや蒸留にかけられる計算資源と時間

シナリオ別の推奨

シナリオ1: エッジデバイスへのデプロイ（スマートフォン、IoT）

メモリとレイテンシの両方が厳しい場合は DistilBERT が第一候補です。パラメータ数66Mは32bit浮動小数点で約264MB、量子化すれば66MB程度まで縮小できます。推論速度もBERT-baseの1.6倍で、リアルタイムアプリケーションに適しています。

精度の低下（GLUE -2.6ポイント）が許容できない場合は、DistilBERTの上にさらにタスク固有の蒸留を行う「二段階蒸留」が有効です。

シナリオ2: メモリが限られたサーバー（多数のモデルを同時稼働）

マイクロサービスアーキテクチャで複数のNLPモデルを同時に稼働させる場合、各モデルのメモリフットプリントが重要になります。この場合は ALBERT-base が適しています。12Mパラメータ（約48MB）と極めて小さく、同じメモリでBERT-baseの9倍のモデルインスタンスを同時にロードできます。

シナリオ3: 精度最優先（研究、コンペティション）

精度が最も重要で計算資源に余裕がある場合は ALBERT-xxlarge が選択肢に入ります。BERT-largeを上回る精度を、より少ないパラメータで達成できます。ただし、推論速度が遅いためリアルタイム用途には不向きです。

シナリオ4: 低レイテンシAPI（検索リランキング、チャットボット）

レイテンシが最優先の場合は DistilBERT です。レイヤー数が半分なので、推論のFLOPsがほぼ半減します。さらに、ONNX RuntimeやTensorRTなどの推論エンジンとの組み合わせ、あるいはINT8量子化を適用することで、さらなる高速化が可能です。

選択のフローチャート

判断基準を整理すると、以下のフローになります。

1. 推論速度が重要か？ → Yes → DistilBERT（またはTinyBERT等のさらなる蒸留モデル）
2. メモリフットプリントが重要か？ → Yes → ALBERT-base
3. 精度を最大化したいか？ → Yes → ALBERT-xxlarge（計算資源があれば）
4. 特に制約がない場合 → BERT-base（最もエコシステムが充実、ファインチューニングの知見が豊富）

ALBERTとDistilBERTは相互排他的ではありません。ALBERTの設計思想（埋め込み分解、パラメータ共有）とDistilBERTの蒸留手法を組み合わせることも可能であり、実際に後続研究ではこのような複合的なアプローチも提案されています。

まとめ

本記事では、BERTの軽量化・効率化手法として ALBERT と DistilBERT を解説しました。

• BERTのパラメータ構成: 約1億1,000万パラメータのうち、Embedding層が21%、Transformer層が77%を占めます。軽量化にはこの2箇所へのアプローチが必要です
• ALBERTの3つの手法: 埋め込み行列の分解（$V \times H \to V \times E + E \times H$）、クロスレイヤーパラメータ共有（12セット → 1セット）、SOPタスク（NSPの改良）により、BERT-baseの約10分の1のパラメータ数で同等の精度を達成します
• DistilBERTの知識蒸留: 温度付きSoftmax、KLダイバージェンス蒸留損失、コサイン埋め込み損失の3つの損失を組み合わせて、BERTの知識を6層モデルに転写します。パラメータ数40%削減・推論1.6倍高速化でBERTの性能の97%を維持します
• 使い分け: メモリ削減にはALBERT、推論速度改善にはDistilBERT。精度最優先ならALBERT-xxlarge、レイテンシ最優先ならDistilBERTが適しています

ALBERTとDistilBERTは、BERTの軽量化という同じ目標に対して、まったく異なる哲学でアプローチしています。ALBERTは「モデル設計の効率化」、DistilBERTは「学習の効率化」です。どちらか一方が優れているのではなく、ユースケースに応じて使い分ける — あるいは組み合わせる — ことが実務では重要です。

BERTの軽量化は現在も活発に研究されており、以下の発展的なトピックも参考にしてください。

次のステップとして、以下の記事も参考にしてください。