Damage-Aware Bandit: новый метод прунинга Transformer-моделей

Суть метода: от жадного выбора к bandit-стратегии

Традиционные методы прунинга часто полагаются на жадные алгоритмы или оценку значимости весов, что может приводить к субоптимальному отбору структурных единиц. Авторы предлагают формулировать задачу отбора функциональных блоков (голов внимания и групп каналов MLP) как проблему Damage-Aware Multi-Armed Bandit. Ключевая инновация — использование paired damage (парного ущерба): разница между потерями на калибровочном батче с замаскированными единицами и базовыми потерями. Это нивелирует вариативность между батчами.

Для выбора единиц применяются два подхода: UCB-style policy и fractional-Beta Thompson Sampling. Отбор происходит последовательно, добавляя по одной единице за шаг, после чего они функционально обнуляются в исходном плотном чекпоинте.

Масштабное тестирование и бенчмарки

Эксперименты охватили широкий спектр архитектур: языковые модели GPT-2, OPT, Pythia, Qwen2.5, SmolLM2 и визуальные модели ViT-B/16, DeiT-Tiny, Swin-Tiny. Оценка проводилась на датасетах WikiText-2, LAMBADA и Imagenette. Метод сравнивался с рандомным отбором, magnitude-based pruning, static-saliency и budgeted-greedy selection.

Статистическая значимость результатов

Результаты демонстрируют статистически значимое превосходство предложенного метода над базовыми стратегиями. Из 28 выделенных сравнений 23 доверительных интервала исключают ноль, а 11 парных тестов имеют p-value < 0.05. После коррекции Бенфамини-Хохберга 6 тестов сохранили значимость (q < 0.05) на полном семействе из 116 тестов.

Почему это важно

Метод позволяет эффективно сжимать модели, сохраняя их функциональность, без необходимости физической компрессии или изменения скорости инференса на этапе отбора. Это открывает путь к более надежному структурированному прунингу для больших языковых и зрительных трансформеров, где сохранение качества критично.

Источник: arXiv cs.AI ↗