Muon Optimizatörü İlk Uzmanların Karışımı (MoE) Katmanını Dengelemeyi Zorlaştırıyor

MoE Mimarilerinde Muon ve Yönlendirici (Router) Dengesi#

Günümüzde birçok açık kaynaklı Mixture-of-Experts (MoE) dil modeli, geleneksel AdamW optimizatörü yerine Muon kullanarak eğitilmektedir. Bu modellerin büyük çoğunluğu, uzmanlar arasındaki iş yükünü dengelemek için her uzmanın yönlendirme (routing) skoruna bir sapma (bias) değeri ekler ve bu değeri her eğitim adımından sonra günceller. Ancak mevcut eğitim reçeteleri, yönlendirici (router) bileşeninin kendisinin de Muon ile eğitilip eğitilmemesi gerektiği konusunda görüş ayrılığı yaşamaktadır.

Deneysel Bulgular ve Sapma Aralığı (Bias Span) Analizi#

55 milyon ve 108 milyon aktif parametreye sahip modeller üzerinde yapılan araştırmada, seçilen optimizatörün bir katmandaki en büyük ve en küçük sapma değeri arasındaki farkı (bias span) nasıl değiştirdiği ölçülmüştür. Tüm token'lar aynı uzmanları tercih ettiğinde bu aralık büyür; dolayısıyla bu metrik bir katmanı dengelemenin ne kadar zor olduğunu doğrudan gösterir:

• İlk Katmanda Belirgin Artış: Daha kısa eğitim koşularında ayarlanan öğrenme oranlarında (learning rate), Muon kullanımı her rastgele tohumlama (seed) değerinde ilk MoE katmanının sapma aralığını yükselterek ortalama iki katına çıkarmıştır. Diğer katmanlarda ise çok daha az değişim gözlemlenmiştir.
• Sorunun Kaynağı: Bu yükseliş, yönlendiriciyi hangi optimizatör eğitirse eğitsin, ağın geri kalanında çalışan Muon optimizatöründen kaynaklanmıştır.
• Ortak Yönelim ve Aşırı Yüklenme: İncelenen koşularda, ilk yönlendiriciye gelen girdilerin ortak bir yön (direction) paylaştığı tespit edilmiştir. Isınma (warmup) aşamasından hemen sonra, bu katmandaki en yoğun uzmanın adil payının dört ila yedi katı arasında yük aldığı görülmüştür.

Çözüm Yöntemleri ve Mühendislik Çıkarımları#

Model eğitimi ve Büyük Model Eğitimi süreçleri açısından araştırmanın sunduğu pratik müdahaleler ve yazılım geliştirme çıkarımları şunlardır:

• Yığın Ortalamasını Çıkarmak (Batch Mean Subtraction): İlk 3.000 adımda, ilk yönlendirici girdilerinden yığın ortalamasını (batch mean) çıkarmak, kayıp (loss) değerini neredeyse hiç değiştirmeden bu aşırı yüklenmeyi ortadan kaldırmıştır. Geliştiriciler için bu, mimariyi değiştirmeden uygulanabilecek düşük maliyetli bir normalizasyon adımıdır.
• Öğrenme Oranı Ayarı: Daha küçük modelde Muon'un öğrenme oranını yarıya indirmek, sapma aralığındaki artışın büyük kısmını gidermiş ve kayıp değerini hafifçe iyileştirmiştir. Kısa koşularda optimize edilen hiperparametrelerin, MoE katmanlarında beklenmedik darboğazlar yaratabileceği unutulmamalıdır.
• Optimizatör Seçimi: Sadece yönlendiriciyi AdamW optimizatörüne taşımak, ısınma süresince yüklerin daha fazla kaymasına neden olmuş ve kaybı biraz artırmıştır. Bu ölçekteki sonuçlar, karmaşık hibrit kurulumlar yerine yönlendiricinin de Muon üzerinde tutulmasını desteklemektedir.

Derin öğrenme mühendisleri ve araştırmacılar, Megatron-LM ile Büyük Model Eğitimi veya özel eğitim boru hatları tasarlarken ilk MoE katmanındaki aktivasyon dağılımlarını yakından izlemelidir. Ağın geri kalanındaki optimizasyon dinamikleri, yönlendirici girdilerinde yönsel çöküşlere yol açarak donanım verimliliğini düşürebilir.