Университет Фудань представляет Lorsa: разбор суперпозиции внимания в трансформерах с помощью разреженных механизмов

Понимание механизмов внимания в трансформерах

Большие языковые модели (LLM) стремительно развиваются, но их внутренние механизмы остаются сложными для понимания. В трансформерах используются множество голов внимания, некоторые из которых имеют конкретные функции, например, индукционные головы, предсказывающие токены по контексту. Однако большинство голов внимания распределяют фокус по разным входам без четко выраженных функций. Такая сложность связана с суперпозицией внимания, когда несколько атомарных единиц внимания накладываются друг на друга в одних и тех же головах, что затрудняет интерпретацию.

Проблемы объяснения работы голов внимания

Ранее исследования выявили специализированные головы внимания с помощью активационного и путевого патчинга, обнаружив функции, такие как перемещение имен, подавление копирования и извлечение длинного контекста. Тем не менее, гипотеза суперпозиции предполагает, что нейроны и головы внимания представляют собой множество перекрывающихся признаков, а не одну функцию. Разреженные автоэнкодеры помогли извлечь разреженные и интерпретируемые признаки из нейросетей, но они все еще сталкиваются с трудностями в объяснении совместного поведения голов внимания в языковых моделях.

Введение в Lorsa: низкоранговое разреженное внимание

Исследовательская команда Университета Фудань представила Lorsa — новый подход для выделения атомарных единиц внимания из сложной суперпозиции в многоголовом самовнимании (MHSA). Lorsa заменяет стандартное MHSA на избыточный набор голов внимания с одномерными OV-схемами и ограничениями разреженности. Такая структура улучшает интерпретируемость, активируя лишь небольшой динамический поднабор голов для каждой позиции токена.

Архитектура и методология Lorsa

Lorsa минимизирует среднеквадратичную ошибку при предсказании выходов MHSA, используя одномерные OV-схемы, ограничивающие операции чтения/записи конкретными признаками остаточного потока. Параметры Query и Key разделяются между головами DLorsa QK, что сохраняет эффективность. В отличие от MHSA, Lorsa активирует только топ-K голов на токен, напоминая Sparse Autoencoders, но с активациями, основанными на паттернах внимания к предыдущим токенам.

Оценка интерпретируемости Lorsa

Команда разработала интерфейс для анализа каждой головы Lorsa, включая топ-активации — токены с наибольшей активацией, и z-паттерн-анализ — разложение вкладов токенов из предыдущих позиций. Эти методы показывают, как работают конкретные головы, например, индукционная голова, фокусирующаяся на токенах "you" и усиливающая соответствующие предсказания.

Результаты и открытия

Lorsa успешно выявил известные механизмы внимания, такие как индукционные, перемещающие имена, последующие и поглощающие головы в моделях Pythia-160M и Llama-3.1-8B. Новые открытия включают арифметические головы, работающие с простыми математическими операциями, и тематические якорные головы, обеспечивающие долгосрочное внимание к тематически связанным токенам, влияя на предсказания в рамках доменной лексики.

Значение и перспективы

Данная работа открывает беспрецедентный взгляд на механизмы внимания в трансформерах и подчеркивает важность учета суперпозиции внимания для интерпретируемости моделей. Несмотря на успехи, остаются задачи полной развязки схем Query-Key и уменьшения эффектов суперпозиции. В будущем планируется исследовать низкоразмерные структуры QK, сквозную суперпозицию и систематическую композицию Query/Key/Value для улучшения понимания и управления языковыми моделями.

Более подробную информацию смотрите в статье, модели на Hugging Face и странице GitHub. Следите за обновлениями в Twitter и исследуйте новости и события об ИИ на Marktechpost.