Сиолошная

Retrieval Head Mechanistically Explains Long-Context Factuality

В архитектуре трансформера есть «головы внимания», на каждом слое их одинаковое количество штук. Эти головы отвечают за то, чтобы перевзвешивать важность слов в контексте: вместо того, чтобы смотреть на тысячи слов за раз модель как бы зануляет большую часть из них, и размазывает 100% внимания лишь по нескольким токенам.

Например, есть голова внимания, которая всегда смотрит только на предыдущее слово, как бы позволяя нейронке опереться на него при генерации. Есть головы, делающие более сложные операции — выявление действующих лиц для местоимений («кошка гуляла, она шла медленно» -> «она» - это кошка). Ничего из этого не программируется вручную — модель во время обучения сама понимает, когда, как и на что смотреть, и изобретает разные алгоритмы.

В рамках науки об интерпретируемости учёные пытаются понять, как учится трансформер, что именно он понимает и как именно работают отдельные механизмы на основе голов внимания. В значимой степени это лишь теория, которая однако имеет огромные перспективы для применения в практике за счёт получения ответов на вопросы в духе «почему модель ошиблась?» или «соврала ли LLM?».

И вот данная работа — как раз такая, она на стыке оптимизации длинного контекста в LLM и интерпретируемости. Для того, чтобы понять текст дальше, нужно прочитать пост <https://www.group-telegram.com/id/seeallochnaya.com/1135> с объяснением принципа теста «иголка в стоге сена».

Авторы придумали критерий, по которому научились определять retrieval heads, которые позволяют копировать модели части промпта. Представьте, что у вас есть 5 страниц текста, и вы в рамках ответа на вопрос приводите цитату: для этого мусолите палец, ставите его на нужное место на странице и двигаете вправо как указатель, и вслед за этим записываете ответ. Вот так и работают эти головы, позволяя модели не сбиваться. Но важно это в очень широком круге задач, особенно в тех, где модель обязана следовать за инструкцией в промпте на 20 страниц — ведь если задуматься это тоже задача поиска: «какое правило тут применить?».

Проанализировав разные семейства моделей (Llama 2, Mistral, Qwen), обнаружили, что такие головы занимают ~5% (~50) от примерно тысячи голов в модели. Их прям конкретно смогли выписать: вот эта вот, с таким то номером. Если их начать отключать (занулять, не давать им смотреть на контекст) — то внезапно модели перестают читать контекст и теряются, их качество существенно падает на задаче поиска иголки в стоге сена. Ещё такой же результат наблюдается на решении математических задач с рассуждением (ведь теперь в них нельзя подсмотреть), а вот качество ответов на вопросы, связанные с чистым знанием не меняется (потому что из контекста не нужно ничего выписывать).

Итак, почему это важно для практики? Когда LLM используется для генерации текста, то для всех предыдущих слов считается огромный тензор вещественных чисел, хранящий информацию о том, что было написано (именно в нём головы внимания находят то, что им нужно). Такой KV-cache для модели LLAMA 2 7B для 100к токенов весит 50 гигабайт. Маленькая модель, контекст не то чтобы очень большой (у Google Gemini вон вообще миллион), и ЦЕЛЫХ 50 ГИГОВ.

В этих гигах как раз таки хранится информация для разных наборов голов. И если мы знаем заранее, что нам нужна голова номер 15 в 10-ом слое — то мы можем сохранять только её, а остальное удалять. Конечно, лучше перестраховаться, и брать, скажем, 10-20% от всего множества — но это в 5-10 раз уменьшает потребление памяти, а значит позволяет и существенно ускорить работу, и уменьшить требования к железу.

Очень жду этой фичи из коробки в большинстве фреймворков, реально очень круто.

www.group-telegram.com/id/seeallochnaya.com/1329

20.8K viewsedited  Apr 25, 2024 at 22:59

Retrieval Head Mechanistically Explains Long-Context Factuality

В архитектуре трансформера есть «головы внимания», на каждом слое их одинаковое количество штук. Эти головы отвечают за то, чтобы перевзвешивать важность слов в контексте: вместо того, чтобы смотреть на тысячи слов за раз модель как бы зануляет большую часть из них, и размазывает 100% внимания лишь по нескольким токенам.

Например, есть голова внимания, которая всегда смотрит только на предыдущее слово, как бы позволяя нейронке опереться на него при генерации. Есть головы, делающие более сложные операции — выявление действующих лиц для местоимений («кошка гуляла, она шла медленно» -> «она» - это кошка). Ничего из этого не программируется вручную — модель во время обучения сама понимает, когда, как и на что смотреть, и изобретает разные алгоритмы.

В рамках науки об интерпретируемости учёные пытаются понять, как учится трансформер, что именно он понимает и как именно работают отдельные механизмы на основе голов внимания. В значимой степени это лишь теория, которая однако имеет огромные перспективы для применения в практике за счёт получения ответов на вопросы в духе «почему модель ошиблась?» или «соврала ли LLM?».

И вот данная работа — как раз такая, она на стыке оптимизации длинного контекста в LLM и интерпретируемости. Для того, чтобы понять текст дальше, нужно прочитать пост <https://www.group-telegram.com/id/seeallochnaya.com/1135> с объяснением принципа теста «иголка в стоге сена».

Авторы придумали критерий, по которому научились определять retrieval heads, которые позволяют копировать модели части промпта. Представьте, что у вас есть 5 страниц текста, и вы в рамках ответа на вопрос приводите цитату: для этого мусолите палец, ставите его на нужное место на странице и двигаете вправо как указатель, и вслед за этим записываете ответ. Вот так и работают эти головы, позволяя модели не сбиваться. Но важно это в очень широком круге задач, особенно в тех, где модель обязана следовать за инструкцией в промпте на 20 страниц — ведь если задуматься это тоже задача поиска: «какое правило тут применить?».

Проанализировав разные семейства моделей (Llama 2, Mistral, Qwen), обнаружили, что такие головы занимают ~5% (~50) от примерно тысячи голов в модели. Их прям конкретно смогли выписать: вот эта вот, с таким то номером. Если их начать отключать (занулять, не давать им смотреть на контекст) — то внезапно модели перестают читать контекст и теряются, их качество существенно падает на задаче поиска иголки в стоге сена. Ещё такой же результат наблюдается на решении математических задач с рассуждением (ведь теперь в них нельзя подсмотреть), а вот качество ответов на вопросы, связанные с чистым знанием не меняется (потому что из контекста не нужно ничего выписывать).

Итак, почему это важно для практики? Когда LLM используется для генерации текста, то для всех предыдущих слов считается огромный тензор вещественных чисел, хранящий информацию о том, что было написано (именно в нём головы внимания находят то, что им нужно). Такой KV-cache для модели LLAMA 2 7B для 100к токенов весит 50 гигабайт. Маленькая модель, контекст не то чтобы очень большой (у Google Gemini вон вообще миллион), и ЦЕЛЫХ 50 ГИГОВ.

В этих гигах как раз таки хранится информация для разных наборов голов. И если мы знаем заранее, что нам нужна голова номер 15 в 10-ом слое — то мы можем сохранять только её, а остальное удалять. Конечно, лучше перестраховаться, и брать, скажем, 10-20% от всего множества — но это в 5-10 раз уменьшает потребление памяти, а значит позволяет и существенно ускорить работу, и уменьшить требования к железу.

Очень жду этой фичи из коробки в большинстве фреймворков, реально очень круто.

BY Сиолошная