論文解説: Muon is Scalable for LLM Training — Newton-Schulz直交化による2倍効率のオプティマイザ

論文概要（Abstract）

Muonオプティマイザは行列直交化に基づく新世代の最適化手法であり、小規模モデルでは優れた性能を示していたが、大規模LLM学習へのスケーラビリティは未検証だった。本論文は、重み減衰（weight decay）の追加とパラメータ別更新スケーリング（per-parameter update scaling）という2つの修正により、Muonをハイパーパラメータチューニングなしで大規模学習に適用可能にした。スケーリング則実験から、MuonはAdamW比で約2倍の計算効率を達成することを示し、その実証として3B/16BパラメータのMoEモデル「Moonlight」を5.7Tトークンで学習した。

この記事は Zenn記事: 2026年版フロンティアLLM学習パイプライン完全解説：事前学習からRLまでの深掘りです。

情報源

arXiv ID: 2502.16982
URL: https://arxiv.org/abs/2502.16982
著者: Jingyuan Liu, Jianlin Su, Zhilin Yang et al.（28名）
発表年: 2025
分野: cs.LG, cs.CL

背景と動機（Background & Motivation）

大規模言語モデル（LLM）の事前学習では、AdamWオプティマイザが事実上の標準として広く使われてきた。AdamWは各パラメータに対して1次・2次モーメントを個別に管理し、適応的な学習率を実現する。しかし、この「パラメータごとの独立した適応」は行列構造の情報を無視しており、更新方向に偏りが生じ得る。Muonオプティマイザは、モメンタム行列をNewton-Schulz反復法で直交化することにより、この偏りを除去する。

従来のMuonは100M規模のモデルで検証されていたが、10B以上のスケールでは (1) 学習不安定性、(2) パラメータ形状の多様性による更新RMSの不均一、(3) 分散環境でのAll-to-All通信オーバーヘッドという3つの課題があった。本論文はこれら3つの課題を系統的に解決する。

主要な貢献（Key Contributions）

貢献1: 重み減衰の追加により大規模学習時の不安定性を解消し、損失スパイクなしの学習を実現
貢献2: パラメータ別更新スケーリング（Adjusted LR法）により、異なる形状の行列間で更新RMSを一定に維持
貢献3: スケーリング則実験でMuonがAdamW比約2倍の計算効率を持つことを定量的に実証
貢献4: 3B/16B MoEモデル「Moonlight」を5.7Tトークンで学習し、同規模モデルのパフォーマンスフロンティアを更新
貢献5: メモリ効率と通信効率を最適化した分散Muon実装をオープンソース公開

技術的詳細（Technical Details）

Newton-Schulz直交化アルゴリズム

Muonの核心は、モメンタム行列 $M_t$ に対するNewton-Schulz反復法による直交化である。目的は $M_t$ の特異値分解 $M_t = U \Sigma V^T$ における直交成分 $O_t = UV^T$ を高速に近似することだ。

初期化として、フロベニウスノルムで正規化する：

\[X_0 = \frac{M_t}{\|M_t\|_F}\]

各反復ステップは以下の多項式で更新する：

\[X_k = a X_{k-1} + b (X_{k-1} X_{k-1}^T) X_{k-1} + c (X_{k-1} X_{k-1}^T)^2 X_{k-1}\]

ここで、

$X_k$: $k$回目の反復後の近似直交行列
$a = 3.4445$, $b = -4.7750$, $c = 2.0315$: 最適化された係数
$X_{k-1} X_{k-1}^T$: グラム行列（対称正定値）

この多項式 $f(x) = ax + bx^3 + cx^5$ は $x=1$ 近傍に不動点を持つよう設計されており、特異値が1に近い成分を素早く収束させる。5回の反復（$N=5$）で十分な精度が得られる。

重み減衰の導入

大規模モデルでは、Muonの更新が特定の方向に蓄積し重みの発散を引き起こす。これを解決するため、パラメータ更新に重み減衰を組み込む：

\[W_t = W_{t-1} - \eta_t (O_t + \lambda W_{t-1})\]

ここで、

$W_t$: 時刻 $t$ のパラメータ行列
$\eta_t$: 学習率
$O_t$: Newton-Schulz直交化後の更新方向
$\lambda$: 重み減衰率（$\lambda = 0.1$）

パラメータ別更新スケーリング（Adjusted LR法）

Muonの直交化後の更新行列 $O_t$ について、$[A, B]$ 形状の行列に対する理論的なRMSは以下で与えられる：

\[\text{RMS}(O_t) = \sqrt{\frac{1}{\max(A, B)}}\]

形状が異なるパラメータ間で更新スケールが不均一になるため、以下の調整を加える：

\[W_t = W_{t-1} - \eta_t \left( 0.2 \cdot O_t \cdot \sqrt{\max(A, B)} + \lambda W_{t-1} \right)\]

スケーリング係数 $0.2$ はAdamWの経験的な更新RMS範囲（0.2〜0.4）に合わせたものである。これにより、すべてのパラメータ行列で更新RMSが約0.2に統一される。

アルゴリズム

以下にMuonオプティマイザの完全な更新ステップを示す：

  
import torch
import torch.distributed as dist

class ScalableMuon:
    """Scalable Muon optimizer with weight decay and adjusted LR.

    Args:
        params: Model parameters
        lr: Learning rate
        momentum: Momentum coefficient (default: 0.95)
        weight_decay: Weight decay ratio (default: 0.1)
        ns_steps: Number of Newton-Schulz iterations (default: 5)
    """

    # Newton-Schulz coefficients (optimized)
    NS_COEFFS = (3.4445, -4.7750, 2.0315)

    def __init__(
        self,
        params: list,
        lr: float = 4.2e-4,
        momentum: float = 0.95,
        weight_decay: float = 0.1,
        ns_steps: int = 5,
    ):
        self.params = list(params)
        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay
        self.ns_steps = ns_steps
        self.state: dict = {}

    def newton_schulz_orthogonalize(self, M: torch.Tensor) -> torch.Tensor:
        """Newton-Schulz iteration for matrix orthogonalization.

        Args:
            M: Momentum matrix of shape (A, B)

        Returns:
            Approximate orthogonal matrix UV^T
        """
        a, b, c = self.NS_COEFFS
        X = M / M.norm()  # Frobenius norm normalization

        for _ in range(self.ns_steps):
            gram = X @ X.T  # Gram matrix (A, A)
            X = a * X + b * gram @ X + c * (gram @ gram) @ X

        return X

    def step(self, lr: float | None = None) -> None:
        """Execute one optimization step.

        Args:
            lr: Optional learning rate override
        """
        current_lr = lr or self.lr

        for p in self.params:
            if p.grad is None:
                continue

            grad = p.grad.data
            if p not in self.state:
                self.state[p] = {"momentum_buffer": torch.zeros_like(grad)}

            buf = self.state[p]["momentum_buffer"]
            buf.mul_(self.momentum).add_(grad)

            if grad.dim() >= 2:
                # Matrix parameter: apply Newton-Schulz orthogonalization
                A, B = grad.shape[0], grad.shape[1]
                O = self.newton_schulz_orthogonalize(buf)
                scale = 0.2 * (max(A, B) ** 0.5)
                update = scale * O + self.weight_decay * p.data
            else:
                # Vector parameter (embedding, bias): fallback to AdamW-style
                update = buf + self.weight_decay * p.data

            p.data.add_(update, alpha=-current_lr)

実装のポイント（Implementation）

分散環境でのメモリ最適化

Muonはモメンタムバッファを1つだけ保持するため、AdamWの2つ（1次・2次モーメント）と比較してメモリオーバーヘッドが50%削減される。分散実装では、勾配のReduce-ScatterとAll-Gatherを組み合わせ、各GPUは自身が担当するパラメータパーティションのみにNewton-Schulz計算を実行する。

通信オーバーヘッド

追加通信コストはAdamWの[1, 1.25]倍に収まる。BF16精度でNewton-Schulz計算を行うことで、FP32比でオーバーヘッドを50%削減できる。実測では前方・後方パス時間の1〜3%程度のレイテンシ増加にとどまる。

Embedding・Head層の扱い

Muonの直交化はEmbedding層やLM Head層には不向きである（次元の非対称性が大きいため）。これらの層にはAdamWを併用するのが標準的なプラクティスだ。Moonlightの学習でもこのハイブリッド戦略が採用されている。

ハイパーパラメータ推奨値

パラメータ	推奨値	備考
学習率（lr）	4.2e-4	WSDスケジュールの最大値
モメンタム	0.95	標準的な値
重み減衰（λ）	0.1	AdamWと同等
NS反復回数	5	精度と速度のバランス
更新RMSスケール	0.2	AdamWの経験的範囲に合わせる

Production Deployment Guide

AWS実装パターン（コスト最適化重視）

Muonオプティマイザを用いたLLM学習パイプラインのAWSデプロイ構成を示す。

トラフィック量別の推奨構成:

規模	月間リクエスト	推奨構成	月額コスト	主要サービス
Small	~3,000 (100/日)	Serverless	$50-150	Lambda + Bedrock + DynamoDB
Medium	~30,000 (1,000/日)	Hybrid	$300-800	Lambda + ECS Fargate + ElastiCache
Large	300,000+ (10,000/日)	Container	$2,000-5,000	EKS + Karpenter + EC2 Spot

Small構成の詳細 (月額$50-150):

Lambda: 1GB RAM, 30秒タイムアウト ($20/月)
Bedrock: Claude 3.5 Haiku, Prompt Caching有効 ($80/月)
DynamoDB: On-Demand ($10/月)
CloudWatch: 基本監視 ($5/月)

Large構成の詳細 (月額$2,000-5,000):

EKS: コントロールプレーン ($72/月)
EC2 Spot (g5.xlarge): 2-4台 (平均$800/月、Spot割引最大90%)
Karpenter: 自動スケーリング（追加コストなし）
Bedrock Batch: 50%割引活用 ($2,000/月)

コスト削減テクニック:

Spot Instances使用で最大90%削減（EKS + Karpenter）
Bedrock Batch API使用で50%削減
Prompt Caching有効化で30-90%削減

コスト試算の注意事項: 上記は2026年2月時点のAWS ap-northeast-1（東京）リージョン料金に基づく概算値です。実際のコストはトラフィックパターンにより変動します。最新料金は AWS料金計算ツールで確認してください。

Terraformインフラコード

Small構成 (Serverless): Lambda + Bedrock + DynamoDB

  
module "vpc" {
  source  = "terraform-aws-modules/vpc/aws"
  version = "~> 5.0"

  name = "muon-llm-vpc"
  cidr = "10.0.0.0/16"
  azs  = ["ap-northeast-1a", "ap-northeast-1c"]
  private_subnets = ["10.0.1.0/24", "10.0.2.0/24"]

  enable_nat_gateway   = false
  enable_dns_hostnames = true
}

resource "aws_iam_role" "lambda_bedrock" {
  name = "muon-lambda-bedrock-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Action = "sts:AssumeRole"
      Effect = "Allow"
      Principal = { Service = "lambda.amazonaws.com" }
    }]
  })
}

resource "aws_iam_role_policy" "bedrock_invoke" {
  role = aws_iam_role.lambda_bedrock.id
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect   = "Allow"
      Action   = ["bedrock:InvokeModel", "bedrock:InvokeModelWithResponseStream"]
      Resource = "arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-5-haiku*"
    }]
  })
}

resource "aws_lambda_function" "llm_handler" {
  filename      = "lambda.zip"
  function_name = "muon-bedrock-handler"
  role          = aws_iam_role.lambda_bedrock.arn
  handler       = "index.handler"
  runtime       = "python3.12"
  timeout       = 60
  memory_size   = 1024

  environment {
    variables = {
      BEDROCK_MODEL_ID    = "anthropic.claude-3-5-haiku-20241022-v1:0"
      DYNAMODB_TABLE      = aws_dynamodb_table.cache.name
      ENABLE_PROMPT_CACHE = "true"
    }
  }
}

resource "aws_dynamodb_table" "cache" {
  name         = "muon-prompt-cache"
  billing_mode = "PAY_PER_REQUEST"
  hash_key     = "prompt_hash"

  attribute {
    name = "prompt_hash"
    type = "S"
  }

  ttl {
    attribute_name = "expire_at"
    enabled        = true
  }
}

Large構成 (Container): EKS + Karpenter + Spot

  
module "eks" {
  source  = "terraform-aws-modules/eks/aws"
  version = "~> 20.0"

  cluster_name    = "muon-training-cluster"
  cluster_version = "1.31"
  vpc_id          = module.vpc.vpc_id
  subnet_ids      = module.vpc.private_subnets

  cluster_endpoint_public_access = true
  enable_cluster_creator_admin_permissions = true
}

resource "kubectl_manifest" "karpenter_provisioner" {
  yaml_body = <<-YAML
    apiVersion: karpenter.sh/v1alpha5
    kind: Provisioner
    metadata:
      name: gpu-spot-provisioner
    spec:
      requirements:
        - key: karpenter.sh/capacity-type
          operator: In
          values: ["spot"]
        - key: node.kubernetes.io/instance-type
          operator: In
          values: ["g5.xlarge", "g5.2xlarge"]
      limits:
        resources:
          cpu: "32"
          memory: "128Gi"
          nvidia.com/gpu: "4"
      ttlSecondsAfterEmpty: 30
  YAML
}

resource "aws_budgets_budget" "monthly" {
  name         = "muon-monthly-budget"
  budget_type  = "COST"
  limit_amount = "5000"
  limit_unit   = "USD"
  time_unit    = "MONTHLY"

  notification {
    comparison_operator        = "GREATER_THAN"
    threshold                  = 80
    threshold_type             = "PERCENTAGE"
    notification_type          = "ACTUAL"
    subscriber_email_addresses = ["ops@example.com"]
  }
}

セキュリティベストプラクティス

IAMロール: 最小権限の原則（Bedrock InvokeModelのみ許可）
ネットワーク: VPC内配置、パブリックサブネット不使用
シークレット: AWS Secrets Manager使用、環境変数ハードコード禁止
暗号化: S3/DynamoDB/EBS全てKMS暗号化
監査: CloudTrail/Config/GuardDuty有効化

運用・監視設定

CloudWatch Logs Insights クエリ:

  
fields @timestamp, model_id, input_tokens, output_tokens
| stats sum(input_tokens + output_tokens) as total_tokens by bin(1h)
| filter total_tokens > 100000

CloudWatch アラーム:

  
import boto3

cloudwatch = boto3.client('cloudwatch')
cloudwatch.put_metric_alarm(
    AlarmName='muon-bedrock-token-spike',
    ComparisonOperator='GreaterThanThreshold',
    EvaluationPeriods=1,
    MetricName='TokenUsage',
    Namespace='AWS/Bedrock',
    Period=3600,
    Statistic='Sum',
    Threshold=500000,
    AlarmActions=['arn:aws:sns:ap-northeast-1:123456789:cost-alerts'],
    AlarmDescription='Bedrockトークン使用量異常'
)

コスト最適化チェックリスト

~100 req/日 → Lambda + Bedrock (Serverless) - $50-150/月
~1000 req/日 → ECS Fargate + Bedrock (Hybrid) - $300-800/月
10000+ req/日 → EKS + Spot Instances (Container) - $2,000-5,000/月
EC2: Spot Instances優先（最大90%削減）
Reserved Instances: 1年コミットで72%削減
Bedrock Batch API: 50%割引（非リアルタイム処理）
Prompt Caching: 30-90%削減
モデル選択: 開発はHaiku($0.25/MTok)、本番はSonnet($3/MTok)
AWS Budgets: 月額予算設定（80%で警告）
CloudWatch: トークン使用量スパイク検知
Cost Anomaly Detection: 自動異常検知
日次コストレポート: SNS/Slackへ自動送信
未使用リソース削除: Lambda Insights活用
タグ戦略: 環境別でコスト可視化
S3ライフサイクル: 古いキャッシュ30日で自動削除
開発環境: 夜間停止設定
Lambda: メモリサイズ最適化
ECS/EKS: アイドルタイムのスケールダウン
Savings Plans: Compute Savings Plans検討
トークン数制限: max_tokens設定で過剰生成防止

実験結果（Results）

スケーリング則実験

4つのモデルサイズ（0.2B, 0.5B, 1.4B, 2.2B）で、異なるトークン量（10B〜200B）を学習し、MuonとAdamWのスケーリング則を比較した。

フィット曲線:

オプティマイザ	スケーリング則
Muon	$L = 2.506 \times C^{-0.052}$
AdamW	$L = 2.608 \times C^{-0.054}$

ここで $L$ は検証損失、$C$ は計算量（FLOPs）である。MuonはAdamWが必要とする計算量の約52%で同等の損失値に到達する。

Moonlightベンチマーク

Moonlight（3B/16B MoE、5.7Tトークン）は、同規模のオープンモデルと比較してパフォーマンスフロンティアを更新した。

ベンチマーク	Moonlight	同規模AdamWモデル	改善
GSM8K	74.4%	68.0%	+6.4%
MATH	42.1%	38.5%	+3.6%
HumanEval	67.1%	62.2%	+4.9%

アブレーション実験

更新RMSの一貫性を保つ手法として3つを比較した。

手法	検証損失	備考
Baseline（0.2·√H スケール）	2.812	不均一なRMS
Update Norm（RMSを0.2に正規化）	2.789	追加計算コスト
Adjusted LR（パラメータ別スケーリング）	2.789	最小コスト

Adjusted LR法は追加計算なしで同等の性能を実現するため、最終的に採用された。

実運用への応用（Practical Applications）

事前学習の高速化

Muonの最大の応用価値は事前学習コストの半減である。例えば、AdamWで1000 GPU時間の学習が必要なモデルを、Muonでは約520 GPU時間で同等品質に到達できる。H100を$3.5/GPU時間で計算すると、1回の学習で$1,680のコスト削減となる。

ハイブリッドオプティマイザ戦略

実運用では、Embedding層とLM HeadにAdamW、それ以外のTransformer層にMuonを使うハイブリッド戦略が推奨される。Moonlightの実装でもこの戦略が採用されており、層ごとのオプティマイザ切り替えはパラメータグループで容易に実装できる。

MoEモデルとの相性

MuonはMoEアーキテクチャと特に相性が良い。MoEの専門家（Expert）パラメータは行列形状が均一であるため、Newton-Schulz直交化の恩恵を最大限に受けられる。MoonlightがMoEアーキテクチャを採用したのもこの理由である。

まとめと今後の展望

Muonオプティマイザは、重み減衰とパラメータ別更新スケーリングという2つのシンプルな修正で大規模LLM学習にスケールする。スケーリング則実験でAdamW比2倍の計算効率を実証し、Moonlight（3B/16B MoE）で実用性を示した。分散環境での通信オーバーヘッドは1〜3%と実用的であり、メモリ使用量もAdamW比50%削減される。

今後は、(1) 100B+パラメータのDenseモデルでの検証、(2) Post-trainingフェーズでのMuon適用可能性、(3) Embedding/Head層の直交化手法の開発が重要な研究方向となる。

参考文献

arXiv: https://arxiv.org/abs/2502.16982
Code: Moonlight pretrained weights and distributed Muon implementation（オープンソース公開）
Related Zenn article: https://zenn.dev/0h_n0/articles/a8792c6407d6e3

📄 論文解説: Muon is Scalable for LLM Training — Newton-Schulz直交化による2倍効率のオプティマイザ

論文概要（Abstract）

情報源

背景と動機（Background & Motivation）

主要な貢献（Key Contributions）

技術的詳細（Technical Details）

Newton-Schulz直交化アルゴリズム

重み減衰の導入

パラメータ別更新スケーリング（Adjusted LR法）

アルゴリズム

実装のポイント（Implementation）

分散環境でのメモリ最適化

通信オーバーヘッド

Embedding・Head層の扱い

ハイパーパラメータ推奨値

Production Deployment Guide

AWS実装パターン（コスト最適化重視）

Terraformインフラコード

セキュリティベストプラクティス

運用・監視設定

コスト最適化チェックリスト

実験結果（Results）

スケーリング則実験

Moonlightベンチマーク

アブレーション実験

実運用への応用（Practical Applications）

事前学習の高速化

ハイブリッドオプティマイザ戦略

MoEモデルとの相性

関連研究（Related Work）

まとめと今後の展望

参考文献

📄 論文解説: Muon is Scalable for LLM Training — Newton-Schulz直交化による2倍効率のオプティマイザ

論文概要（Abstract）

情報源

背景と動機（Background & Motivation）

主要な貢献（Key Contributions）

技術的詳細（Technical Details）

Newton-Schulz直交化アルゴリズム

重み減衰の導入

パラメータ別更新スケーリング（Adjusted LR法）

アルゴリズム

実装のポイント（Implementation）

分散環境でのメモリ最適化

通信オーバーヘッド

Embedding・Head層の扱い

ハイパーパラメータ推奨値

Production Deployment Guide

AWS実装パターン（コスト最適化重視）

Terraformインフラコード

セキュリティベストプラクティス

運用・監視設定

コスト最適化チェックリスト

実験結果（Results）

スケーリング則実験

Moonlightベンチマーク

アブレーション実験

実運用への応用（Practical Applications）

事前学習の高速化

ハイブリッドオプティマイザ戦略

MoEモデルとの相性

関連研究（Related Work）

まとめと今後の展望

参考文献

関連記事

📄 論文解説: Muon + MLA + MoE — 3技術統合で68%メモリ削減・3.2倍推論高速化を実現

📄 論文解説: Kimi K2 — 1兆パラメータMoE+MuonClipで実現するオープンエージェント知能

📄 論文解説: MegaScale-MoE — 1,440 GPU上で1.88倍高速化を実現するMoE学習システム