Каждому параметру задается несколько возможных опций разной стоимости
c_k, а также распределения вероятности p_i = softmax(\alpha_i) (\alpha_i - некие исходные логиты). И у каждого параметра есть свой вес w_i. Итоговая стоимость имеет вид:C = \sum_i w_i <p_i, c>Фиксированный бюджет C задает некоторое гладкое многообразие в пространстве логитов \alpha, потому можно применять известные техники Римановой оптимизации для нахождения решения.
Алгоритм работает следующим образом:
🔅 Проецируем градиент на многообразие
🔅 Делаем шаг любым оптимизатором (например, Adam-ом)
🔅 Проецируем полученную точку обратно на многообразие
🔅 Проецируем первый момент на текущую касательную плоскость
Логиты зашумляются случайным шумом. Сила шума постепенно понижается от некоторой максимальной до минимальной. Тем самым в начале активно исследуем пространство решений, а затем аккуратно сходимся.
В качестве лосс-функции используется KL дивергенция между исходной и сжатой моделями.
🧪 Эксперименты
Метод сначала валидируют на задаче с рюкзаком, где есть точный алгоритм динамического программирования. Методы с Лагранжианом / расширенным Лагранжианом не находят целевое решение, а Риманова оптимизация находит.
Затем пробуют сжимать экспертов в MoE (OLMoE, Qwen3-30B-A3B, Qwen3-Coder-Next). Заметно опережают по качеству REAP и EvoESAP (эволюционный алгоритм) будучи значительно быстрее последнего.
Однако, метод чувствителен к выбору данных, на Qwen3-Coder если обучаться на коде, хорошо сохраняется качество на коде, но просаживается на остальных доменах. И наоборот, калибровочные данные общего вида просаживают сильно качество на коде.
Для mixed precision квантизации показывают себя не хуже эволюционных алгоритмов, будучи в несколько раз дешевле по стоимости. При 2.5-3.5 битной квантизации опережают MxMoE бейзлайн на 1% в среднем по качеству.
💡 Выводы
Как будто бы сравнительно недорогая и рабочая техника для структурированного прунинга / неравномерной квантизации. Стоит попробовать!