Budget Constraints as Riemannian Manifolds

以下は、Michael Helcig と Dan Alistarh による論文「Budget Constraints as Riemannian Manifolds」の詳細な技術的概要です。

1. 問題定義

本論文は、効率的な機械学習において繰り返し直面する課題である予算制約付き離散割り当て問題を取り扱います。目的は、モデル損失（例）などのグローバルな目的関数を最小化しつつ、モデルサイズやレイテンシなどの総コスト予算の制約の下で、各グループ（例：ニューラルネットワークの層）に $K$ 個の離散オプション（例：量子化ビット幅、プルーニングの決定、エキスパートの選択）のいずれかを割り当てることです。

主要な課題:

• 非分解可能な目的関数: モデル損失はすべての割り当てに依存して決まり（例：層間の相互作用）、独立したグループごとの項の和として記述できません。これにより、標準的な動的計画法（DP）は実際の損失を直接最適化できず、DP は層ごとの感度などの「代理」目的関数のみを最適化できます。
• 既存手法の限界:
  • 組合せソルバー（DP）: 予算を正確に扱いますが、分解可能な目的関数を必要とするため、不正確な代理スコアの使用を余儀なくされます。
  • 進化探索: 実際の損失を評価しますが、勾配を持たないため、数百回のフォワードパスが必要となり、高い計算コストに苦しみます。
  • ペナルティ/ラグランジュ法: 損失にペナルティ項（例：$\lambda(C-B)^2$）を追加します。これらは予算を近似してのみ強制し、しばしば振動し、感度の高いハイパーパラメータ調整を必要とします。

2. 手法：リーマン制約最適化（RCO）

著者らは、予算制約をペナルティとしてではなく、ログit空間における滑らかなリーマン多様体として扱う幾何学的アプローチを提案します。

A. 予算多様体

著者らは、各グループ $i$ の割り当て分布をログit $\alpha_i \in \mathbb{R}^K$ でパラメータ化し、確率を $p_i = \text{softmax}(\alpha_i)$ と定義します。総期待コストは以下のように定義されます：
$$C(\alpha) = \sum_{i=1}^N w_i \langle p_i, c \rangle$$
ここで、$w_i$ はグループの重み、$c$ はコストベクトルです。

• 多様体の構造: レベルセット $\mathcal{M} = \{ \alpha \mid C(\alpha) = B \}$ は、滑らかな $(NK-1)$ 次元のリーマン部分多様体であることが証明されています。
• 幾何学的性質:
  1. 閉形式の法線: 勾配 $\nabla C(\alpha)$ は閉形式で利用可能であり、（異なるコストを仮定すれば）ゼロではなく、多様体に対する法線ベクトルを定義します。
  2. 単調な再帰（Retraction）: コストベクトルに沿ってログitをシフトさせると、期待コストが単調に変化します。これにより、アルゴリズムは複雑な反復射影ではなく、単純な二分探索（またはニュートン法）を用いて任意の点を多様体 $\mathcal{M}$ 上に射影できます。
  3. 安価なベクトル輸送: 一つの接平面から別の接平面へ運動量を移動させる操作は、単一の内積演算に帰着します。

B. RCO アルゴリズム

提案されたアルゴリズムは、標準的な 1 次最適化（Adam）を 3 つの幾何学的操作でラップし、各ステップで予算を正確に強制します：

1. 接空間への射影: 損失勾配を、法線ベクトル $\nabla C(\alpha)$ 方向の成分を除去することで多様体の接空間 onto 射影します。これにより、オプティマイザは予算を維持する方向のみを認識します。
2. 再帰（Retraction）: Adam ステップ後、反復点が多様体から外れる可能性があります。アルゴリズムはすべてのログitをコストベクトルに沿ってシフトさせ（$\alpha' = \alpha + t \cdot c$）、$C(\alpha') = B$ となる $t$ を二分探索で求めます。
3. 運動量輸送: オプティマイザの運動量（1 次モーメント）を新しい接空間 onto 射影して整合性を維持します。

C. 離散実行可能性の処理（フォワードパス）

多様体は連続的なログit空間で動作するため、本論文は RCO と以下を組み合わせます：

• Gumbel-STE（Straight-Through Estimator）: フォワードパスで離散的な割り当てを生成するために使用します。
• 予算制約付き動的計画法: フォワードパス中に離散割り当て問題（多選択ナップサック問題）を正確に解き、評価される構成が予算を満たすことを保証します。
• メカニズム: フォワードパスでは、摂動されたログitを与えられた場合に最適な離散構成 $z^*$ を見つけるために DP を使用します。バックワードパスでは、摂動されたログitの softmax 確率を通じて勾配を流すために STE を使用し、多様体射影がこれらの勾配が予算制約に違反しないことを保証します。

3. 主要な貢献

1. 予算多様体の特定: 著者らは、softmax 期待コストのレベルセットが、固有の計算安価な幾何学的性質（閉形式の法線、単調な再帰）を持つリーマン多様体を形成することを特定しました。
2. RCO アルゴリズム: ゼロの制約ハイパーパラメータ（調整する $\lambda$ なし）で正確な予算制約を強制する新しい最適化フレームワークです。接空間射影、二分探索による再帰、および Adam 周回での運動量輸送を組み合わせています。
3. 不等式制約の処理: 多様体をスラック変数 $s$ を通じて不等式制約（$C(\alpha) \le B$）に拡張し、オプティマイザが未使用の予算を劣ったオプションに浪費することなく自然に活用できるようにします。
4. 離散 - 連続ハイブリッド: Gumbel-STE と予算制約付き DP の統合により、正確な離散実行可能性を維持しながら、実際の非分解可能なモデル損失の最適化を可能にします。

4. 実験結果

本論文は、合成および実世界の大型言語モデル（LLM）圧縮タスクにおいて RCO を検証しました。

• 合成ナップサック問題（MCKP）:

  • 困難なインスタンスにおいて、RCO は動的計画法（DP）の最適解を回復します。
  • ラグランジュ法は最適値の**83%**で頭打ちとなり、顕著な制約違反（$10^{-1}$ から $10^{-2}$）を示しますが、RCO は機械精度に近い違反（$10^{-9}$）を維持します。
  • RCO はスラック変数を利用することで「安価な最適解」（最良の解が予算より低いコストの場合）を完璧に処理しますが、等式制約付き手法は予算を浪費します。

• MoE エキスパートのプルーニング:

  • OLMoE-1B-7B および Qwen3-30B-A3B で評価されました。
  • RCO は、複数のベンチマーク（ARC、MMLU など）において、EvoESAP（進化探索のベースライン）の精度と同等かそれ以上を達成します。
  • 効率性: RCO は、これらの結果を3〜16 倍低いウォールクロック時間で達成します（例：EvoESAP の 5.8 時間に対し約 10 分）。

• 混合精度量子化:

  • Qwen3-8B（252 層にビット幅 2〜8 を割り当て）で評価されました。
  • 高圧縮（2.25 ビット）において、RCO は HIGGS（線形代理手法）と比較してパープレキシティを**36%削減し、EvoPress（進化探索）と比較して4%**削減します。
  • RCO は 2.5 ビットで EvoPress と同じ最適値に収束しますが、約 10 倍短い時間（11〜14 時間に対し 62 分）で達成します。

5. 意義と影響

• 正確な制約強制: ペナルティ法とは異なり、RCO は各反復で予算が正確に満たされることを保証し、ペナルティ重みのハイパーパラメータ調整の必要性を排除します。
• 非分解可能損失の最適化: 離散実行可能性（フォワードパスで DP により処理）と連続最適化（バックワードパスで多様体により処理）を分離することで、RCO は厳格な予算制約の下で複雑な非分解可能なモデル損失（エンドツーエンドの精度など）の直接最適化を可能にします。
• 計算効率: 幾何学的操作（射影、再帰、輸送）は計算的に安価（本質的に内積と二分探索）であり、進化探索が過度に高価となる大規模モデルにもスケーラブルです。
• 汎用性: このフレームワークは、線形コストを伴う softmax パラメータ化を含む任意の問題に適用可能であり、モデル圧縮を超えて、リソース制約付きニューラルアーキテクチャ探索やその他の離散割り当て問題へ拡張できます。

要約すると、本論文は予算制約付き離散最適化を、特定のリーマン多様体をナビゲートする問題として再定義し、ペナルティベースおよび進化アプローチに対する、理論的根拠があり、ハイパーパラメータ不要で、極めて効率的な代替手段を提供しています。