О механизме внимания
Трансформер-модели обладают интересным свойством: во время генерации текста они могут обращаться к информации из любого предыдущего токена. Но при этом нужно быть избирательным — если каждый токен будет влиять на следующий одинаково, система окажется неэффективной. Здесь нужен механизм, который решает, с какой силой один токен влияет на другой.
Оказывается, этот механизм можно визуализировать!
При наведении или нажатии на любой сгенерированный токен на экране появляются предыдущие токены, которые повлияли на его создание.
О точности визуализации
Стоит уточнить: термин «повлиял» здесь упрощен. Визуализация сильно редуцирует информацию — она вычисляет вес внимания, масштабирует его по величине вектора значений, агрегирует данные со всех attention heads и суммирует по всем слоям модели. Полученное значение управляет прозрачностью отображения предыдущих токенов. Наибольшие значения всегда отображаются непрозрачно, остальные интерполируются между ними.
Чтобы свести сложную информацию к одному числу на токен, пришлось отбросить многое. При разработке автор даже сомневался, будет ли результат понятен. Однако инструмент выявил интересные закономерности.
Примеры из практики
В примере по умолчанию («Office Move Summary») хорошо видно, как работает копирование. При наведении на дословно скопированный текст — адреса, даты — исходные данные явно выделяются, потому что сгенерированный токен черпает большую часть информации из исходного источника.
Это объясняет кажущийся парадокс: если модели предсказывают токены вероятностно, почему они так хорошо копируют данные? Разве они не должны в конце концов ошибиться просто по статистическим причинам?
Но механизм внимания разрешает эту проблему. Модель не восстанавливает последовательность из ограниченного внутреннего состояния — у неё есть полный доступ ко всем прошлым токенам. Она просто выбирает, из каких токенов брать информацию при копировании. Поэтому вероятность ошибок остаётся очень низкой. В примере «Debugging an Average Function» видно, как небольшая модель (600 млн параметров) с лёгкостью воспроизводит всю функцию на JavaScript, кроме самого исправления. (Хотя самостоятельно найти ошибку она не может, ей нужны подсказки.)
Ещё один любопытный случай — слово «remain» в фразе «Existing access cards and phone numbers remain» в примере про переезд офиса. При наведении на него видно, что модель обращается одновременно к слову «work» в фразе «Existing employee access cards will work» и к выражению «stay the same» в «company phone numbers will stay the same». Модель как бы комбинирует информацию из обеих фраз — это выглядит довольно элегантно.
Реализация
Сама визуализация — это достаточно простое React-приложение с использованием Transformers.js для генерации текста. Но чтобы вытащить из модели дополнительные данные для визуализации, пришлось отойти от стандартного цикла генерации. Разработчик переписал генерационный цикл непосредственно в приложении, чтобы отслеживать значения для визуализации.
Размер модели исчисляется сотнями мегабайт, и просить пользователя ждать, пока она загрузится, просто неприемлемо. Поэтому были предварительно сгенерированы примеры запросов, которые загружаются и отображаются моментально.
Другая сложность в том, что часть данных в визуализации не является стандартным выходом модели. Если бы разработка велась на Python с библиотеками для machine learning, это было бы просто. Но Transformers.js использует файлы .onnx, содержащие весь граф вычислений. Логика загрузки модели и вычислений реализована на WebAssembly, поэтому нет простого способа получить доступ к чему-то, кроме предопределённых выходов.
В итоге был написан небольшой скрипт, который модифицирует файл onnx и открывает доступ к внутренним значениям. Это означает, что нельзя просто использовать стандартную модель .onnx. Для работы в браузере потребовалось загрузить инструментированную версию модели в собственный репозиторий на Hugging Face и указать приложению на неё.
Весь код доступен в репозитории на GitHub.