予算制約をリーマン多様体として扱う

Abstract

総コストの予算のもとで、N個のグループそれぞれにK個の選択肢のうち1つを割り当てることは、機械学習における反復的な問題であり、混合精度量子化、非一様プルーニング、エキスパート選択に現れます。目的（モデル損失）はすべての割当てに対して同時に依存しており、グループごとに分解できないため、組合せソルバが真の目的を直接最適化することを妨げ、代理目的に限られてしまいます。進化的探索は実際の損失を評価できますが勾配情報を欠いており、一方、罰則ベースの手法は勾配を与えますが予算を厳密には満たさず、さらに感度の高いハイパーパラメータ調整が必要です。私たちは、softmax緩和のもとでは、予算制約がロジット空間において滑らかなリーマン多様体を定義し、特に幾何学が単純になることを観察します：法線ベクトルは閉形式で得られ、コストベクトルに沿ってロジットをシフトさせると期待コストが単調に変化するため、二分探索によるリトラクションが可能になり、ベクトル輸送は単一の内積に帰着します。この構造に基づき、標準的なAdam更新を、接空間への射影、二分探索によるリトラクション、モメンタム輸送で拡張するリーマン拘束最適化（RCO）を提案します。Gumbelのstraight-through推定と、離散的な実行可能性のための予算制約付き動的計画法を組み合わせることで、RCOは制約ハイパーパラメータを導入することなく、厳密な予算遵守のもとで真の目的を一次（勾配）に基づいて最適化できるようになります。最適解が既知の合成ナップサック問題では、多様体に基づく制約処理により最適解が復元されますが、罰則法は最適の83%で頭打ちになります。LLM圧縮タスク、すなわち混合精度量子化やMoEエキスパートプルーニングでは、RCOは進化的探索手法に対して一致または上回る性能を示し、評価された構成において壁時計コストを3倍〜16倍低く抑えます。