TokenFormer: Переосмысление масштабирования трансформера с токенизированной моделью Параметры

TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

October 30, 2024

Авторы: Haiyang Wang, Yue Fan, Muhammad Ferjad Naeem, Yongqin Xian, Jan Eric Lenssen, Liwei Wang, Federico Tombari, Bernt Schiele

cs.AI

Аннотация

Трансформеры стали преобладающей архитектурой в основных моделях благодаря их отличной производительности в различных областях. Однако значительная стоимость масштабирования этих моделей остается значительной проблемой. Эта проблема возникает в основном из-за их зависимости от фиксированного количества параметров в пределах линейных проекций. При внесении архитектурных изменений (например, размеров каналов) весьма часто требуется переобучение всей модели с нуля. Поскольку размеры моделей продолжают увеличиваться, такой подход приводит к все более высоким вычислительным затратам и становится неустойчивым. Для преодоления этой проблемы мы представляем TokenFormer, архитектуру, которая масштабируется собственным образом и использует механизм внимания не только для вычислений между входными токенами, но также для взаимодействий между токенами и параметрами модели, тем самым улучшая архитектурную гибкость. Обрабатывая параметры модели как токены, мы заменяем все линейные проекции в Трансформерах нашим слоем внимания к токенам-параметрам, где входные токены действуют как запросы, а параметры модели - как ключи и значения. Эта переформулировка позволяет постепенное и эффективное масштабирование без необходимости переобучения с нуля. Наша модель масштабируется от 124 миллионов до 1,4 миллиарда параметров путем пошагового добавления новых пар ключ-значение параметров, достигая производительности, сравнимой с Трансформерами, обученными с нуля, при значительном снижении затрат на обучение. Код и модели доступны по адресу https://github.com/Haiyang-W/TokenFormer.

English

Transformers have become the predominant architecture in foundation models due to their excellent performance across various domains. However, the substantial cost of scaling these models remains a significant concern. This problem arises primarily from their dependence on a fixed number of parameters within linear projections. When architectural modifications (e.g., channel dimensions) are introduced, the entire model typically requires retraining from scratch. As model sizes continue growing, this strategy results in increasingly high computational costs and becomes unsustainable. To overcome this problem, we introduce TokenFormer, a natively scalable architecture that leverages the attention mechanism not only for computations among input tokens but also for interactions between tokens and model parameters, thereby enhancing architectural flexibility. By treating model parameters as tokens, we replace all the linear projections in Transformers with our token-parameter attention layer, where input tokens act as queries and model parameters as keys and values. This reformulation allows for progressive and efficient scaling without necessitating retraining from scratch. Our model scales from 124M to 1.4B parameters by incrementally adding new key-value parameter pairs, achieving performance comparable to Transformers trained from scratch while greatly reducing training costs. Code and models are available at https://github.com/Haiyang-W/TokenFormer.