Share Your Attention: Transformer Weight Sharing via Matrix-based Dictionary Learning

## КОНТЕКСТ И ПРОБЛЕМАТИКА Large language models (LLMs) и Vision Transformers (ViT) заложили основу для большого количества передовых приложений в области искусственного интеллекта. Однако их внедрение в реальные системы часто сталкивается с существенными препятствиями, связанными с высокими вычислительными и памятными требованиями. Для снижения этих требований были разработаны различные методы компрессии, такие как низкоранговая аппроксимация, удаление голов внимания и кэширование ключей-значений (KV). Однако большинство этих методов сосредоточено на оптимизации внутри блоков (intra-block), не уделяя должного внимания возможностям оптимизации между блоками (inter-block). Повторяющаяся структура трансформеров, состоящая из многократно повторяющихся слоев, подразумевает высокую степень избыточности между этими слоями. Несмотря на то, что существуют некоторые подходы к использованию этой избыточности, такие как кэширование ключей-значений, этот вопрос остается недостаточно изученным. При этом, методы компрессии, основанные на извлечении статистических регулярностей между слоями, могли бы существенно снизить количество параметров без существенного ущерба для качества. Вдохновленные подходами диктонического обучения (dictionary learning) в сверточных нейронных сетях (CNNs), авторы предлагают новый подход для структурированного общения весов (weight sharing) между слоями трансформеров. Основная идея заключается в том, чтобы разделить проекционные матрицы внимания (attention projection matrices) на общие атомы (shared dictionary atoms), что позволит существенно уменьшить количество параметров, необходимых для модели. ## ПРЕДЛОЖЕННЫЙ МЕТОД Предлагаемый метод, называемый MASA (Matrix Atom Sharing in Attention), представляет собой оптимизацию структуры весовых матриц в трансформерах. Он разработан для сокращения избыточности в параметрах модели, особенно в модулях внимания (attention modules). Метод основывается на диктоническом обучении, где веса представляются в виде линейных комбинаций общих атомов (shared dictionary atoms). В рамках MASA, проекционные матрицы внимания разделяются на общие компоненты, которые затем используются для представления весов каждого слоя. Это позволяет уменьшить количество необходимых параметров на 66,7%, при этом сохраняя качество модели на уровне состояния искусства. Особенностью MASA является то, что он может быть использован как drop-in replacement для стандартных трансформеров, не требуя сложных изменений в архитектуре или дополнительных процессов, таких как дистилляция модели. Метод обучается с помощью стандартных оптимизаторов, что делает его легко встраиваемым в существующие модели. Он также позволяет эффективно захватывать статистические регулярности между слоями, что делает его особенно эффективным в задачах, где необходим высокий уровень обобщения. ## ЭКСПЕРИМЕНТАЛЬНЫЕ РЕЗУЛЬТАТЫ Авторы провели исследования на моделях различных размеров, от 100 миллионов до 700 миллионов параметров. Результаты экспериментов показали, что MASA достигает высокого качества на различных бенчмарках, превосходя существующие базовые подходы, такие как grouped-query attention (GQA), низкоранговые базисные методы и недавно предложенные подходы Repeat-all-over/Sequential sharing. В частности, MASA демонстрирует лучшие результаты по точности и перплексии по сравнению с этими методами при сравнительном количестве параметров. Аблационные исследования показали, что метод является достаточно робастным к размеру словаря (dictionary size) и эффективен в захвате статистических регулярностей между слоями. Кроме того, MASA был успешно применен к Vision Transformers (ViT) для задач классификации изображений и детекции. Результаты показали, что MASA может сохранить качество модели на уровне современных подходов, снизив количество параметров на 66,7%. ## ПРАКТИЧЕСКАЯ ЗНАЧИМОСТЬ Предлагаемый метод MASA имеет значительный практический потенциал в области развертывания моделей трансформеров в реальных приложениях. Благодаря существенному сокращению количества параметров, MASA позволяет уменьшить вычислительные и памятные требования моделей, что делает их более доступными для развертывания на устройствах с ограниченными ресурсами. Кроме того, MASA может быть применен для оптимизации предобученных LLMs, позволяя сократить их размер без существенного ухудшения качества. Это может быть особенно полезно в задачах, требующих быстрого и эффективного выполнения моделей на устройствах с ограниченными вычислительными мощностями. ## ВЫВОДЫ И ПЕРСПЕКТИВЫ В рамках данной работы был предложен новый метод MASA для эффективного общения весов между слоями трансформеров. Этот метод показал свою эффективность на различных задачах, включая классификацию текста и изображений. Будущие исследования могут быть направлены на дальнейшее улучшение метода, включая его применение к более крупным моделям и задачам, а также исследование возможностей интеграции MASA с другими методами компрессии.