Сканирование и фиксация: изучение динамики обучения и композиции токенов в однослойном трансформере

Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer

May 25, 2023

Авторы: Yuandong Tian, Yiping Wang, Beidi Chen, Simon Du

cs.AI

Аннотация

Архитектура Transformer продемонстрировала впечатляющие результаты в различных исследовательских областях и стала основой многих моделей нейронных сетей. Однако понимание того, как она работает, остается ограниченным. В частности, остается загадкой, как представление формируется в процессе градиентного обучения при использовании простой функции потерь для предсказания. В данной работе мы анализируем динамику обучения с использованием стохастического градиентного спуска (SGD) для однослойного Transformer, состоящего из одного слоя самовнимания и одного декодирующего слоя, в задаче предсказания следующего токена, придерживаясь математической строгости. Мы раскрываем "черный ящик" динамического процесса, в котором слой самовнимания комбинирует входные токены, и выявляем природу лежащей в основе индуктивной предвзятости. Более конкретно, при допущениях (а) отсутствия позиционного кодирования, (б) длинной входной последовательности и (в) более быстрого обучения декодирующего слоя по сравнению со слоем самовнимания, мы доказываем, что самовнимание действует как дискриминативный алгоритм сканирования: начиная с равномерного внимания, оно постепенно уделяет больше внимания уникальным ключевым токенам для предсказания конкретного следующего токена и меньше внимания общим ключевым токенам, которые встречаются для разных следующих токенов. Среди уникальных токенов оно постепенно снижает веса внимания, следуя порядку от низкой до высокой совместной встречаемости ключевого и запрашиваемого токенов в обучающей выборке. Интересно, что этот процесс не приводит к ситуации "победитель получает все", а замедляется из-за фазового перехода, который контролируется скоростями обучения двух слоев, оставляя (почти) фиксированную комбинацию токенов. Мы проверяем эту динамику "сканирования и фиксации" на синтетических и реальных данных (WikiText).

English

Transformer architecture has shown impressive performance in multiple research domains and has become the backbone of many neural network models. However, there is limited understanding on how it works. In particular, with a simple predictive loss, how the representation emerges from the gradient training dynamics remains a mystery. In this paper, for 1-layer transformer with one self-attention layer plus one decoder layer, we analyze its SGD training dynamics for the task of next token prediction in a mathematically rigorous manner. We open the black box of the dynamic process of how the self-attention layer combines input tokens, and reveal the nature of underlying inductive bias. More specifically, with the assumption (a) no positional encoding, (b) long input sequence, and (c) the decoder layer learns faster than the self-attention layer, we prove that self-attention acts as a discriminative scanning algorithm: starting from uniform attention, it gradually attends more to distinct key tokens for a specific next token to be predicted, and pays less attention to common key tokens that occur across different next tokens. Among distinct tokens, it progressively drops attention weights, following the order of low to high co-occurrence between the key and the query token in the training set. Interestingly, this procedure does not lead to winner-takes-all, but decelerates due to a phase transition that is controllable by the learning rates of the two layers, leaving (almost) fixed token combination. We verify this \emph{scan and snap} dynamics on synthetic and real-world data (WikiText).