Как устроен Representation Finetuning, родившийся из идеи интерпретируемости LLM

Black box AI is bad AI —  гласит слоган исследовательской группы Pr (AI)2R (Practical AI Alignment and Interpretability Research). Её основал прошлым летом стэнфордский автор Аттикус Гигер (Atticus Geiger). Своей миссией группа считать превратить AI в «хороший AI»,   то есть сделать его интерпретируемым.  

Пока авторы выпустили три работы: Rigorously Assessing Natural Language Explanations of Neurons (лучшая статья 2023 по версии BlackBoxNLP), в которой попытались провести интерпретацию на уровне нейронов, Linear Representations of Sentiment in Large Language Models,   где исследовали репрезентацию настроения в LLM и RAVEL: Evaluating Interpretability Methods on Disentangling Language Model Representations, где представили бенчмарк для оценки интерпретируемости. Есть и более ранние работы Гигера, в частности, он предложил исследовать внутренности LLM с помощью интервенций (изменения внутренних состояний). Суть проста: если зафиксировать скрытое состояние, и выход модели поменяется так, как будто какой-либо компонент производил это состояние, то это даёт нам право установить причинно-следственную связь.   Но тут расскажем о том, к каким конструктивным идеям приводит исследование интерпретируемости. Как говорится, критикуешь — предлагай.  

Работы по интерпретируемости LLM приводят к выводу — скрытые состояния трансформера скрывают в себе много семантической информации. Именно они, а не веса. Опираясь на это авторы из Стэнфорда, в том числе и Гигер, пришли к гипотезе — донастройка модели должна модифицировать не веса, а скрытые состояния. Да, обновление небольшой доли весов LLM действительно позволяет эффективно донастроить её под нужную задачу, и эта мысль породила успешные и ставшие классическими LoRA и DoRA, но если учесть интерпретируемость, то результаты оказываются ещё лучше. 

Новый метод авторы назвали по аналогии с PEFT (Parameter-efficient finetuning) — ReFT, Representation Finetuning (статья). Он основан на тех самых интервенциях, которые предложил Гигер. Модифицировать будем не веса, а прямо самые скрытые состояния. Интервенция (I) состоит из трёх компонентов — собственно функции интервенции (Ф) и своего рода координат L и P. Первая «координата» указывает на то, какой слой скрытый состояний мы изменяем, а вторая — какие из токенов. 

На примере с картинки выше интервенция касается только репрезентаций двух токенов (1 и 3) со слоя 2. Другие состояния при этом не затрагиваются. Это тоже важное отличие от PEFT методов, в которых обновляются все внутренние состояния. Ф забирает значение состояния сразу же, как только оно посчитано и возвращает в то же место. То есть на следующий слой подаётся уже измененное состояние. 

Этот подход определяет сразу целое семейство методов. Один из возможных методов реализовали сами авторы в той же работе. Это Low-Rank Linear Subspace ReFT (LoReFT). Функция интервенции в этом случае содержит три обучаемых параметра — две матрицы R и W, и вектор b

В экспериментах с LoReFT изменяли только состояния для двух первых и двух последних токенов. Остальные не трогали. В блоге авторы признаются, что вообще-то не ожидали большого успеха от метода. Интерес был скорее теоретический, с точки зрения интерпретируемости. А вышло, что LoReFT обошла текущий state-of-the-art 


Что же из этого следует? Во-первых, обнаружилось, что интервенция во выходные токены влияет на токены всех поколений. То есть можно контролировать все поколения, поместив несколько начальных токенов в нужные состояния. Например, на картинке сверху на графике для instruct-tuning, LoReFT обошла текущего лидера, с помощью интервенции в 4 токена 4 первых слоев. Во-вторых, ReFT даёт новые инсайты для дальнейшего изучения интерпретируемости. Пока авторы делают вывод, что возможно сопоставлять нейроны каким-то словам вообще не имеет смысла. Возможно, у них нет никакой «специализации», точнее есть, но меняется с каждым новым входным запросом. Что еще более важно, то, что кодируют нейроны, зависит от исходных вычислений, в которых они участвуют. 

Почему LoReFT работает? Иронично, но авторы пока не очень  понимают, снова получается black box. Но надеются понять из экспериментов — изучать численные пределы и из этого попробовать выяснить секрет действия. Например, авторы взяли один промпт на английском и обучали интервенции восстанавливать начало Алисы в стране чудес. То есть пытались проверить сколько слов может удержать одна интервенция, изменяя количество слоёв и токенов. Оказалось —  до 2048 слов. 


Чтобы легче было перейти с PEFT на ReFT, авторы создали библиотеку pyreft — можно пробовать.

Больше наших обзоров AI‑статей — на канале Pro AI.

© Habrahabr.ru