PyTorch を使用した Mixture of Experts (MoE) モデルの実装

Mixtral 8x7B の発売は、オープン AI の分野、特に Mixture-of-Experts (MoEs) の概念で幅広い注目を集めています。専門家のミックス (MoE) コンセプトは、協調的知性の象徴であり、「全体は部分の総和よりも大きい」という言葉を体現しています。 MoE モデルは、さまざまな専門家モデルの長所を組み合わせて、より優れた予測を提供します。これは、ゲーティング ネットワークと一連のエキスパート ネットワークを中心に構築されており、各ネットワークは特定のタスクの異なる側面で優れています。

この記事では、Pytorch を使用して MoE モデルを実装します。コードに入る前に、Mixture of Experts のアーキテクチャを簡単に紹介しましょう。

MoEアーキテクチャ

MoEは（1）エキスパートネットワークと（2）ゲーティングネットワークの2種類のネットワークで構成されています。

エキスパート ネットワーク: エキスパート ネットワークは独自のモデルであり、それぞれがデータのサブセットで適切に機能するようにトレーニングされています。 MoE の考え方は、問題領域を包括的にカバーするために、補完的な強みを持つ複数の専門家を配置することです。

ゲート ネットワーク: ゲート ネットワークは、個々の専門家からの貢献の指揮者、コーディネーター、または管理者として機能します。どのネットワークがどのタイプの入力を処理するのに適しているかを学習（または重み付け）します。トレーニングされたゲーティング ネットワークは、新しい入力ベクトルを評価し、熟練度に基づいて最も適切な専門家または専門家の組み合わせに処理責任を割り当てることができます。ゲーティング ネットワークは、エキスパートの出力と現在の入力の関連性に基づいて重みを動的に調整し、カスタマイズされた応答を保証します。

上図はMoEにおける処理フローを示しています。専門家混合モデルの利点は、そのシンプルさにあります。 MoE モデルは、複雑な問題空間と、問題を解決する際に専門家がどのように反応するかを学習することで、単独の専門家よりも優れたソリューションを生み出すのに役立ちます。ゲートネットワークは効果的なマネージャーとして機能し、状況を評価して、最適な専門家にタスクを渡します。新しいデータが入ってくると、モデルは新しい入力に対する専門家の強みを再評価することで適応し、柔軟な学習アプローチを実現します。

MoE は機械学習モデルの導入に多大なメリットをもたらします。注目すべき利点が 2 つあります。

文部科学省の強みは、専門家ネットワークの多様性と専門性にあります。 MoE 設定は、単一のモデルでは達成が難しい精度で多面的な問題を処理できます。

MoE は本質的にスケーラブルです。タスクの複雑さが増すにつれて、他の専門家のモデルを変更することなく、より多くの専門家をシステムにシームレスに統合できるようになり、専門知識の範囲が拡大します。つまり、MoE は、事前にトレーニングされた専門家を機械学習システムにパッケージ化するのに役立ちます。

Mixture of Experts モデルは、レコメンデーション システム、言語モデリング、さまざまな複雑な予測タスクなど、多くの分野で応用されています。 GPT-4 は複数の専門家で構成されているという噂があります。確認はできませんが、GPT-4 のようなモデルは、MoE アプローチを通じて複数のモデルのパワーを活用することで、最良の結果を提供します。

Pytorch コード

ここでは、Mixtral 8x7B などの大規模モデルで使用される MOE テクノロジについては説明しません。代わりに、あらゆるタスクに適用できるシンプルなカスタム MOE を作成します。コードを通じて、MOE の動作原理を理解することができます。これは、大規模モデルで MOE がどのように機能するかを理解するのに非常に役立ちます。

以下では、PyTorch コードの実装をセクションごとに紹介します。

ライブラリをインポートします。

 import torch import torch.nn as nn import torch.optim as optim

エキスパートモデルを定義します。

 class Expert(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim): super(Expert, self).__init__() self.layer1 = nn.Linear(input_dim, hidden_dim) self.layer2 = nn.Linear(hidden_dim, output_dim) def forward(self, x): x = torch.relu(self.layer1(x)) return torch.softmax(self.layer2(x), dim=1)

ここでは、単純なエキスパート モデルを定義します。これは 2 層 MLP であり、relu アクティベーションを使用し、最後に softmax を使用して分類確率を出力することがわかります。

ゲーティング モデルを定義します。

 # Define the gating model class Gating(nn.Module): def __init__(self, input_dim, num_experts, dropout_rate=0.1): super(Gating, self).__init__() # Layers self.layer1 = nn.Linear(input_dim, 128) self.dropout1 = nn.Dropout(dropout_rate) self.layer2 = nn.Linear(128, 256) self.leaky_relu1 = nn.LeakyReLU() self.dropout2 = nn.Dropout(dropout_rate) self.layer3 = nn.Linear(256, 128) self.leaky_relu2 = nn.LeakyReLU() self.dropout3 = nn.Dropout(dropout_rate) self.layer4 = nn.Linear(128, num_experts) def forward(self, x): x = torch.relu(self.layer1(x)) x = self.dropout1(x) x = self.layer2(x) x = self.leaky_relu1(x) x = self.dropout2(x) x = self.layer3(x) x = self.leaky_relu2(x) x = self.dropout3(x) return torch.softmax(self.layer4(x), dim=1)

ゲートモデルはより複雑で、過剰適合を防ぐための正規化のための 3 つの線形層とドロップアウト層を備えています。 ReLU および LeakyReLU 活性化関数を使用して非線形性を導入します。最後の層の出力サイズはエキスパートの数に等しく、これらの出力にソフトマックス関数が適用されます。エキスパートの出力と組み合わせることができるように重みを出力します。

注: 実際、ゲーティング ネットワーク、またはルーティング ネットワークは、エキスパート モデルへの入力を制御するため、MOE の最も複雑な部分であり、ゲーティング ネットワークには多くの設計スキームがあります。たとえば (私の記憶が正しければ)、Mixtral 8x7B は 8 人のエキスパートのうち上位 2 人だけを採用します。したがって、ここではさまざまなソリューションについて詳しく説明するのではなく、その基本原則とコード実装のみを紹介します。

完全なMOEモデル:

 class MoE(nn.Module): def __init__(self, trained_experts): super(MoE, self).__init__() self.experts = nn.ModuleList(trained_experts) num_experts = len(trained_experts) # Assuming all experts have the same input dimension input_dim = trained_experts[0].layer1.in_features self.gating = Gating(input_dim, num_experts) def forward(self, x): # Get the weights from the gating network weights = self.gating(x) # Calculate the expert outputs outputs = torch.stack( [expert(x) for expert in self.experts], dim=2) # Adjust the weights tensor shape to match the expert outputs weights = weights.unsqueeze(1).expand_as(outputs) # Multiply the expert outputs with the weights and # sum along the third dimension return torch.sum(outputs * weights, dim=2)

ここでは主に、入力と各エキスパートの出力の予測を通じて重みを計算し、最終的に重みを使用してすべてのエキスパートの結果を合計し、モデルの出力を取得する順方向伝播コードについて説明します。

これは「アンサンブル学習」に少し似ていますか?

テスト

実装の簡単なテストをしてみましょう。まず、簡単なデータセットを生成します。

 # Generate the dataset num_samples = 5000 input_dim = 4 hidden_dim = 32 # Generate equal numbers of labels 0, 1, and 2 y_data = torch.cat([ torch.zeros(num_samples // 3), torch.ones(num_samples // 3), torch.full((num_samples - 2 * (num_samples // 3),), 2) # Filling the remaining to ensure exact num_samples ]).long() # Biasing the data based on the labels x_data = torch.randn(num_samples, input_dim) for i in range(num_samples): if y_data[i] == 0: x_data[i, 0] += 1 # Making x[0] more positive elif y_data[i] == 1: x_data[i, 1] -= 1 # Making x[1] more negative elif y_data[i] == 2: x_data[i, 0] -= 1 # Making x[0] more negative # Shuffle the data to randomize the order indices = torch.randperm(num_samples) x_data = x_data[indices] y_data = y_data[indices] # Verify the label distribution y_data.bincount() # Shuffle the data to ensure x_data and y_data remain aligned shuffled_indices = torch.randperm(num_samples) x_data = x_data[shuffled_indices] y_data = y_data[shuffled_indices] # Splitting data for training individual experts # Use the first half samples for training individual experts x_train_experts = x_data[:int(num_samples/2)] y_train_experts = y_data[:int(num_samples/2)] mask_expert1 = (y_train_experts == 0) | (y_train_experts == 1) mask_expert2 = (y_train_experts == 1) | (y_train_experts == 2) mask_expert3 = (y_train_experts == 0) | (y_train_experts == 2) # Select an almost equal number of samples for each expert num_samples_per_expert = \ min(mask_expert1.sum(), mask_expert2.sum(), mask_expert3.sum()) x_expert1 = x_train_experts[mask_expert1][:num_samples_per_expert] y_expert1 = y_train_experts[mask_expert1][:num_samples_per_expert] x_expert2 = x_train_experts[mask_expert2][:num_samples_per_expert] y_expert2 = y_train_experts[mask_expert2][:num_samples_per_expert] x_expert3 = x_train_experts[mask_expert3][:num_samples_per_expert] y_expert3 = y_train_experts[mask_expert3][:num_samples_per_expert] # Splitting the next half samples for training MoE model and for testing x_remaining = x_data[int(num_samples/2)+1:] y_remaining = y_data[int(num_samples/2)+1:] split = int(0.8 * len(x_remaining)) x_train_moe = x_remaining[:split] y_train_moe = y_remaining[:split] x_test = x_remaining[split:] y_test = y_remaining[split:] print(x_train_moe.shape,"\n", x_test.shape,"\n", x_expert1.shape,"\n", x_expert2.shape,"\n", x_expert3.shape)

このコードは、0、1、2 の 3 つのクラス ラベルを持つ合成データセットを作成します。クラス ラベルに基づいて機能を操作し、モデルが学習できるデータに何らかの構造を導入します。

データは、個々の専門家のトレーニング セット、MoE モデル、およびテスト セットに分割されます。エキスパート モデルがサブセットでトレーニングされるようにします。これにより、最初のエキスパートはラベル 0 と 1 で十分にトレーニングされ、2 番目のエキスパートはラベル 1 と 2 でより適切にトレーニングされ、3 番目のエキスパートはラベル 2 と 0 をより多く参照するようになります。

予想される結果は、ラベル 0、1、2 の各専門家の分類精度は満足できるものではないものの、3 人の専門家の決定を組み合わせることで、MoE のパフォーマンスが向上するというものです。

モデルの初期化とトレーニング設定:

 # Define hidden dimension output_dim = 3 hidden_dim = 32 epochs = 500 learning_rate = 0.001 # Instantiate the experts expert1 = Expert(input_dim, hidden_dim, output_dim) expert2 = Expert(input_dim, hidden_dim, output_dim) expert3 = Expert(input_dim, hidden_dim, output_dim) # Set up loss criterion = nn.CrossEntropyLoss() # Optimizers for experts optimizer_expert1 = optim.Adam(expert1.parameters(), lr=learning_rate) optimizer_expert2 = optim.Adam(expert2.parameters(), lr=learning_rate) optimizer_expert3 = optim.Adam(expert3.parameters(), lr=learning_rate)

エキスパートモデルとMoEモデルがインスタンス化されます。トレーニング損失を計算するための損失関数を定義し、トレーニング中に重みの更新を実行するために各モデルのオプティマイザーを設定します。

トレーニングの手順も非常に簡単です

# Training loop for expert 1 for epoch in range(epochs): optimizer_expert1.zero_grad() outputs_expert1 = expert1(x_expert1) loss_expert1 = criterion(outputs_expert1, y_expert1) loss_expert1.backward() optimizer_expert1.step() # Training loop for expert 2 for epoch in range(epochs): optimizer_expert2.zero_grad() outputs_expert2 = expert2(x_expert2) loss_expert2 = criterion(outputs_expert2, y_expert2) loss_expert2.backward() optimizer_expert2.step() # Training loop for expert 3 for epoch in range(epochs): optimizer_expert3.zero_grad() outputs_expert3 = expert3(x_expert3) loss_expert3 = criterion(outputs_expert3, y_expert3) loss_expert3.backward()

各エキスパートは、基本的なトレーニング ループを使用して、データの異なるサブセットで個別にトレーニングされます。ループは指定されたエポック数だけ繰り返されます。

以下はMOEのトレーニングです

# Create the MoE model with the trained experts moe_model = MoE([expert1, expert2, expert3]) # Train the MoE model optimizer_moe = optim.Adam(moe_model.parameters(), lr=learning_rate) for epoch in range(epochs): optimizer_moe.zero_grad() outputs_moe = moe_model(x_train_moe) loss_moe = criterion(outputs_moe, y_train_moe) loss_moe.backward() optimizer_moe.step()

MoE モデルは、事前にトレーニングを受けた専門家によって作成され、その後別のデータセットでトレーニングされます。トレーニング プロセスは単一のエキスパートの場合と似ていますが、ゲーティング ネットワークの重みはトレーニング プロセス中に更新されます。

最後に評価関数は次のようになります。

 # Evaluate all models def evaluate(model, x, y): with torch.no_grad(): outputs = model(x) _, predicted = torch.max(outputs, 1) correct = (predicted == y).sum().item() accuracy = correct / len(y) return accuracy

評価関数は、指定されたデータに対するモデルの精度を計算します (x はサンプルを表し、y は予想されるラベルを表します)。精度は、正しい予測の数と予測の総数の比率として計算されました。

結果は次のとおりです。

 accuracy_expert1 = evaluate(expert1, x_test, y_test) accuracy_expert2 = evaluate(expert2, x_test, y_test) accuracy_expert3 = evaluate(expert3, x_test, y_test) accuracy_moe = evaluate(moe_model, x_test, y_test) print("Expert 1 Accuracy:", accuracy_expert1) print("Expert 2 Accuracy:", accuracy_expert2) print("Expert 3 Accuracy:", accuracy_expert3) print("Mixture of Experts Accuracy:", accuracy_moe) #Expert 1 Accuracy: 0.466 #Expert 2 Accuracy: 0.496 #Expert 3 Accuracy: 0.378 #Mixture of Experts Accuracy: 0.614

見ることができます

エキスパート 1 は、テスト データセット内のサンプルの約 46.6% のクラス ラベルを正しく予測します。

エキスパート 2 のパフォーマンスはわずかに優れており、正しい予測率は約 49.6% でした。

エキスパート 3 は 3 人のエキスパートの中で最も精度が低く、サンプルの約 37.8% が正しく予測されました。

MoE モデルは各専門家を大幅に上回り、全体的な精度は約 61.4% でした。

要約する

私たちのテストの出力結果は、専門家混合モデルの威力を示しています。このモデルは、ゲーティング ネットワークを通じてさまざまなエキスパート モデルの利点を組み合わせ、単一のエキスパート モデルよりも高い精度を実現します。ゲーティング ネットワークは、入力データに基づいて各専門家の貢献度を評価する方法を効果的に学習し、より正確な予測を生成します。ハイブリッド エキスパートは、個々のモデルのさまざまな専門知識を活用し、テスト データセットでより優れたパフォーマンスを提供します。

また、既存のタスクのテストに MOE を使用して、より良い結果を得ることができることも示しています。