Стягивай куда нужно: Activation Steering Tutorial

Привет, друзья! Если вы по запросу «как сделать модель добрее» видите в output-е LLM фразу «рулевое управление» — значит LLM говорит про Steering. В этом туториале вы:

узнаете, что такое steering и на чем он основан;
осуществите steering, используя pytorch-hooks;
познакомитесь с библиотеками nnsight и pyvene для interventions;

И если какое-то слово из bullet-ов было непонятно, они все станут вам понятны к концу.

Activation Steering — это

В research-народье, Activation Steering — это добавление, вычитание или иная трансформация векторов во внутренних состояниях LLM во время forward pass-а. Steering основан на предположении о том, что у обученной модели есть фиксированные «направления» в латентном пространстве.

Activation Steering — это inference-time intervention (вмешательство в модель во время инференса). Мы не меняем веса модели (в отличие от fine-tuning) — мы вмешиваемся в поток вычислений «на лету», пока модель «думает» — то есть генерирует.

Базовая формула:
$\text{hidden\_state}^{(\ell)} \;\leftarrow\; \text{hidden\_state}^{(\ell)} + \alpha \cdot \mathbf{v}$

где $\mathbf{v}$ — steering vector, вектор, кодирующий нужное поведение, $\ell$ — номер слоя, в который мы вмешиваемся, $\alpha$ — сила вмешательства.

Сдвигаемое поведение должно быть чётко выражено и иметь полярную пару, например:

refusal vs compliance;
positive sentiment vs negative sentiment;

Заметим, что во втором случае «positive» и «negative» определяется контекстом. Классический пример из жизни — то, что «positive» для консервативных людей, явно «negative» для сторонников нового. Отложим это пока в памяти.

В этом туториале мы поставим цель сдвинуть модель в сторону hate-speech. Выбор темы hate-speech обусловлен исследовательским интересом. Сдвигать, повторюсь, можно в любое место, выражающее полярность.

Примеры из ноутбука не выражают мою личную позицию относительно субъектов высказывания.

Что нужно для steering?

Первое — конечно, модель. Для быстрого демо используется небольшая модель gpt2, чтобы ноутбук запускался почти везде.

Скрытый текст

Для более веселых экспериментов можно заменить MODEL_NAME в ноуктбуке, который я прикреплю ниже на:

— gpt2-medium
— EleutherAI/pythia-410m
— TinyLlama/TinyLlama-1.1B-Chat-v1.0
— Llama/Mistral/Gemma open-weight модели, если есть доступ и GPU

Выбирайте своё!

Contrastive dataset

Первый шаг стиринга — конструирование направления. Чтобы его найти, нам нужен набор данных, отражающий сдвигаемую полярность. Поскольку наша цель hate-speech, рассмотрим mixed_hate_dataset, где каждое высказывание имеет одну из двух меток:

0 (is_harmfull_opposition): ненавистническое / дискриминационное высказывание. Пример:

«Mentally retarded people are uneducated and should not be accepted into schools.»
1 (is_harmfull_opposition): опровержение / tolerant

«Mentally retarded people can be educated and should be accepted into schools.»

Датасет собран так, что каждое harmfull имеет safe пару. Поэтому он хорош для steering. Идеален он был бы, если бы все топики были из одной темы, но лучшее — враг хорошего и его нам достаточно.

В такой постановке данных, мы ожидаем, что steering vector будет указывать из пространства ненавистнических высказываний в пространство высказываний толерантных:

steering_vector = mean(acts_tolerant) − mean(acts_hate)

По постановке вектора, применяя его с alpha > 0, мы толкаем генерацию в сторону tolerant (потому что весь hate из tolerant мы вычли). С alpha < 0 — в обратную сторону. У этого есть нюансы и их вы увидите ниже.

Для получения направления мы будем использовать residual stream модели на выбранном слое.

Residual stream

GPT-2, как и большинство трансформеров, устроен по принципу residual connections: каждый блок не «обрабатывает» тензор с нуля, а добавляет свой вклад к уже существующему:

$\mathbf{x}^{(\ell+1)} = \mathbf{x}^{(\ell)} + \text{Block}_\ell\bigl(\mathbf{x}^{(\ell)}\bigr)$

Это значит, что один и тот же «поток» — residual stream — идет от входа до выхода через все слои. Каждый слой читает из него и дописывает в него. model.transformer.h[layer] — это выход residual block $\ell$ , то есть $\mathbf{x}^{(\ell+1)}$ после применения слоя. Он же именуется как hidden_state.

Почему residual stream важен для нас:

информация в нём накапливается аддитивно — если не получилось на слое «до» мы можем понадеятся, что концепт просто живет в другом слое;
информация по определению вынуждена читаться линейно, отсюда линейного сдвига нам достаточно (во многом так как механизм внимания собран целиком из линейных проекций: $Q = W_Q \mathbf{x}, K = W_K \mathbf{x}, V = W_V \mathbf{x}$ )

Представление последнего токена.

Заметим, что в residual stream, на самом деле, много токенов. Мы будем брать последний.

GPT-2 — авторегрессионная модель с causal attention: токен на позиции видит только токены с позициями $\leq i$ . Это значит, что последний токен видит весь предшествующий контекст и является естественным «сборщиком» информации о промпте.

Альтернативы существуют — усреднение по всем токенам, взвешенное по attention — но они непонятны для интерпретации и почти не используются.

Hook.

Hook — это функция-перехватчик, которую вы «вешаете» на модуль. Она вызывается автоматически при каждом forward pass:

handle = model.transformer.h[layer].register_forward_hook(hook_fn)# hook_fn(module, input, output) — вызывается после вычисления слоя

PyTorch вызовет hook_fn сразу после того, как слой завершит вычисление. Можно:

читать output и извлекать активации (как здесь)
возвращать модифицированный тензор и делать steering

Технический момент — после работы — обязательно handle.remove(), иначе hook останется висеть на модели навсегда. Пример хука — ниже. Если я накосячила с отступами — простите, но у вас будет тетрадь.

Построение steering vector

Метод называется Contrastive Activation Addition (CAA) — из статьи Steering Llama 2 via Contrastive Activation Addition. Идея:

Берём два набора промптов: позитивный класс (tolerant) и негативный (hate).
Для каждого промпта снимаем активацию последнего токена на слое $\ell$.
Вычисляем разность средних:

$\mathbf{v} = \underbrace{\frac{1}{|P^+|} \sum_{i \in P^+} \mathbf{h}_i}_{\text{mean(tolerant)}} \;-\; \underbrace{\frac{1}{|P^-|} \sum_{j \in P^-} \mathbf{h}_j}_{\text{mean(hate})}$
Нормализуем: $\hat{\mathbf{v}} = \mathbf{v} / \|\mathbf{v}\|_2$

Стоп. Ведь пару абзацев выше ты сказала, что среднее бессмысленно.

Среднее по токенам внутри примера — не то, потому что смешивает разные вычислительные роли.

Среднее же по примерам — статистика над понятием — если понятие (например, «токсичность») кодируется в residual stream, то оно соответствует некоторому направлению $\mathbf{v} \in \mathbb{R}^d$ .

Тогда активация последнего токена для -го примера раскладывается как: $\mathbf{x}i = \mathbf{x}\text{base} + \alpha_i \mathbf{v} + \boldsymbol{\varepsilon}_i$

где $\alpha_i = +1$ для позитивных примеров, $\alpha_i = -1$ для негативных, $\boldsymbol{\varepsilon}_i$ — шум.

Тогда:

$\bar{\mathbf{x}}^+ - \bar{\mathbf{x}}^- = 2\mathbf{v} + \underbrace{(\bar{\boldsymbol{\varepsilon}}^+ - \bar{\boldsymbol{\varepsilon}}^-)}_{\to 0 \text{ при } N \to \infty}$

На константу — забили.

Почему нормализуем?

Без нормализации длина $\mathbf{v}$ зависит от:

количества промптов в датасете
«разброса» активаций в конкретном слое
масштаба самой модели

После нормализации $\hat{\mathbf{v}}$ — единичный вектор, и $\alpha$ становится единственным параметром, управляющим силой вмешательства.

def build_steering_vector(pos_prompts: List[str], neg_prompts: List[str], layer: int) -> torch.Tensor:    pos = get_block_output_activations(pos_prompts, layer)    neg = get_block_output_activations(neg_prompts, layer)    vec = (pos - neg).mean(dim=0)    vec = vec / (vec.norm() + 1e-8)    return vec.to(DEVICE)STEER_LAYER = 6steering_vector = build_steering_vector(positive_prompts, negative_prompts, STEER_LAYER)print("vector shape:", steering_vector.shape)print("norm:", steering_vector.norm().item()) # Подумайте, почему норма равна единице. # Если вы не дали ответ сходу — вернитесь к теории!

Sanity check

Чтобы понять, сто сделанный вектор осмысленнее, можно использовать проекцию в семантику методом logit lens. Про него я писала (но для картинок), для текстов то же самое — мы пропускаем вектор через lm_head и смотрим, какие токены он «предпочитает» и «избегает».

То, что мы получим, не будет точно — вектор находится в пространстве активаций residual stream, а не финальном logit-пространстве, но даст интуицию о семантическом содержании вектора — и отсюда может быть полезно как sanity check. На подопытном датасете, мы с вами получим следующее:

Список против получился вполне себе hate, однако список с оложительными словами вышел шумным — например, мусор для нашей задачи — точно токены stellar, NAV, incorpor, eele. Очевидно, что мы нашли не идеальный вектор, он кодирует ось «агрессивная лексика vs что-то ещё», а не строго «hate vs tolerant».

Кроме того, набор данных имеет небольшой шум — он содержит некоторое количество фраз, где оппозит — лишь отрицание (например, «Blacks are not dangerous»). Такие предложения содержат hate-тему как субъект и модель может активировать на них hate-направление. Разность средних следовательно зашумлена по определению.

Однако, несмотря на все ограничения, направление hate speech мы имеем, и всё равно можем использовать его для практики. Однако учтем не идеальную полярность. И на этом месте самое время достать нашу отложку из памяти! Помните, что «positive» и «negative» определяется контекстом? Вот вам и классическое gargbage in = garbage out.

Стирим

Как работает интервенция

На каждом decode-шаге (генерация одного нового токена) модель делает полный forward pass. Наш hook перехватывает выход слоя $\ell$ и модифицирует его:
$\mathbf{x}^{(\ell+1)} \;\leftarrow\; \mathbf{x}^{(\ell+1)} + \alpha \cdot \hat{\mathbf{v}}$

Дальше слои $\ell+1, \ell+2, \ldots$ продолжают вычисления уже с изменённым тензором. При генерации новых токенов вмешательство повторяется при каждом decode-шаге — модель постоянно находится под воздействием вектора.

Почему она работает?

Ответ даёт наше предположение — если понятие «tolerant» кодируется как линейное направление $\hat{\mathbf{v}}$ , то добавление $\alpha \hat{\mathbf{v}}$ к активации буквально перемещает внутреннее состояние модели в ту часть пространства, которая ассоциируется с выбранной полярностью.

Слои, стоящие выше $\ell$ , «видят» смещённое состояние и продолжают обработку как будто этот контекст изначально был в нужню-сторону-ориентированным. Это работает, потому что трансформеры с residual connections обрабатывают информацию аддитивно — каждый слой читает из общего потока и пишет в него. Наша добавка не «ломает» вычисления — она смещает точку отсчёта.

А куда добавлять?

Мы можем добавить как на все токены, так и только на последний. На практике характеристики такие:

last: сдвигаем только позицию последнего токена — минимальное вмешательство, меньше побочных эффектов, влияет на следующий предсказанный токен через attention дальше по сети
all: сдвигаем все позиции — более агрессивно, меняет весь контекст, который attention будет читать на следующих слоях

Для baseline-экспериментов обычно берут all — эффект сильнее и легче заметен. Но в ноутбуке используется last.

Теперь точно стирим.

Если вы запустите ноутбук (а я горячо (именно это слово!) рекомендую это), то увидите следующую картинку.

Вектор оказался инвертирован: alpha > 0 двигает модель в сторону hate, а не tolerant. Это один из типичных failure mode контрастивных датасетов с минимальными парами — tolerant-примеры лексически похожи на hate (те же субъекты, та же тема), и модель активирует hate-направление на обоих классах. Разность средних тогда указывает не туда, куда ожидалось.

Не баг — фича.

Для математической согласованности мы могли бы инвертировать вектор, но мы можем сделать проще, и договориться, что для найденного вектора:

— alpha < 0 двигает модель в сторону hate,
— alpha > 0 — в сторону tolerant.

Знак подобран эмпирически через тест на генерации, а не из логики датасета в силу неидеального контраста. Это норм.

Logit lens, кстати, этого не предсказал — он показал hate-токены на минус-стороне, то есть выглядел «правильно». Это напоминание, что logit lens — аппроксимация: он проецирует вектор напрямую через lm_head, игнорируя то, что с возмущением сделают последующие слои. Эмпирический тест на генерации надёжнее. Мы просто инвертируем вектор и двигаемся дальше.

Остался открытый вопрос: какое значение alpha выбрать для старта? Слишком маленькое — эффект незаметен, слишком большое — модель теряет когерентность (способность генерировать связный текст).

Откуда брать init $\alpha$ ?

Достанем наше знание о том, что steering vector нормирован: $|\hat{\mathbf{v}}| = 1$ . Откуда $\alpha$ — это количество единиц, на которое мы сдвигаем активацию вдоль направления концепта. Сдвиг осмыслен только относительно того, насколько велики сами активации в этом слое. Возможная отправная точка — взять среднюю норму активаций на выбранном слое:
$\alpha_\text{init} \approx \mathbb{E}\bigl[|\mathbf{h}^{(\ell)}|\bigr]$

и брать сколько-то от нее. Например, начать с 10% от типичной нормы и увеличивать вдвое до тех пор, пока не появится эффект. На практике для GPT-2 это означает $\alpha \in [20, 100]$ , для больших моделей — другой масштаб. Слишком большой $\alpha$ даёт oversteering: модель теряет когерентность (свою способность вообще норм отвечать) и начинает повторяться или выдавать бессмыслицу — это сигнал, что вы вышли за пределы тренировочного распределения активаций.

Перебирая константны, мы обнаруживаем, что достаточно $\alpha = 16$ (≈20% от нормы активаций на слое), чтобы модель согласилась с нашим утверждением. Но оценивать по одному примеру некорректно — один промпт может быть нетипичным, а эффект зависеть от конкретной фразы, а не от вектора в целом. Поэтому переходим к количественному eval: прогоним steering по набору примеров и измерим, как меняется доля «yes»/»no» ответов в зависимости от $\alpha$ . Бенчмарк вы можете найти в ноутбуке, я лишь остановлюсь на мотивации и ограничениях выбора дизайна.

Почему yes/no?

Это максимально простой дизайн: мы принуждаем модель выдать один из двух сигналов. Это убирает вариативность языка — не надо парсить открытый текст и гадать, «поддержала» ли модель утверждение.

Ограничения есть:

GPT-2 не instruction-tuned, поэтому отвечает на вопросы непредсказуемо — часть ответов не содержит ни yes, ни no
вопрос-форма меняет распределение токенов относительно train distribution модели

Для серьёзного eval нужен toxicity classifier (например unitary/toxic-bert) или LLM-as-judge — они дают более надёжную оценку без зависимости от yes/no поведения модели. Но мы ограничились keyword baseline.

Оптимальный $\alpha$ на выбранном слое — 20. При нем модель согласна на 42% hate утверждениях (против 22% при baseline), и отрицает только 27% (против 73% на baseline) и реже соглашается с позитивными утверждениями (32% при baseline против 9% со стирингом).

Часть 2. Если вы здесь — вы молодец!

К этому моменту мы построили steering через сырые PyTorch hooks — минимальный, но уже сильный подход. Steering вкусный, steering популярный — и вокруг него есть библиотеки. Поэтому в части B (сейчас) посмотрим на две из них: nnsight и pyvene — они делают то же самое чище. Заодно убедимся, что все три реализации дают одинаковый результат.

NNsight

nnsight — библиотека для интерпретации и интервенций во внутренности deep learning моделей. Она позволяет читать и изменять активации через tracing context.

Deferred execution — ключевая идея

Главное отличие от ручных hooks — отложенное выполнение. Синтаксис:

with nn_model.trace(prompt):    acts = nn_model.transformer.h[layer].output[:, -1, :].save()

В таком стиле вы не читаете активацию немедленно. Вы описываете план: «при следующем forward pass — сохрани это». nnsight строит граф операций, выполняет forward pass, и только потом .save() возвращает реальный тензор.

Преимущество этого: не нужно регистрировать и снимать hooks вручную — нет риска «забыть снять hook» и сломать модель для всех следующих вызовов. Код читается декларативно: «я хочу видеть x» вместо «перехвати слой, положи в кэш, сними».

То же самое для интервенций:

with nn_model.trace(prompt):    hidden = nn_model.transformer.h[layer].output    nn_model.transformer.h[layer].output[:] = hidden + alpha * vector

nnsight автоматически применяет это вмешательство в нужный момент forward pass.

API nnsight развивается между версиями. Этот раздел написан для nnsight >= 0.6; если что-то не работает — проверьте nnsight.__version__ и документацию.

Все можно запустить, в теле статьи посмотрим лишь на главный блок. Steering выглядит так — в tracing context можно заменить / модифицировать активацию модуля.

with nn_model.generate(prompt, max_new_tokens=...):    hidden = nn_model.transformer.h[layer].output[0]    hidden[:, :, :] = hidden + alpha * vector    output = nn_model.generator.output.save()

В разных версиях nnsight генерация и сохранение output могут немного отличаться. При запуске ответы будут те же, с точностью до токена. Кроме того, для вектора, который получите через NNsight, у вас будет такой вывод:

nn_vector shape: torch.Size([768]),  norm: 1.0000Cosine similarity (nnsight vs HF hooks): 1.0000(Close to 1.0 = identical vectors; different APIs produce the same result)

А значит, этому миру кровавого open-source можно доверять!

Pyvene

pyvene (Stanford NLP) заменяет ручные hooks декларативным конфигом: вы описываете что менять, а не как.

Идея: intervention как объект

В ручном подходе интервенция — это функция-hook. В pyvene — это объект с типом и конфигом:

config = IntervenableConfig([{  "layer": L,  "component": "block_output",          # где перехватывать  "intervention_type": AdditionIntervention,  # что делать   }])pv_model = IntervenableModel(config, model)

IntervenableModel оборачивает оригинальную модель и автоматически управляет hooks — вам не нужно регистрировать и снимать их вручную.

AdditionIntervention — что происходит внутри

При вызове pv_model(inputs, unit_locations=..., source_representations=...) в указанных позициях выполняется:
$\mathbf{h} \;\leftarrow\; \mathbf{h} + \mathbf{s}$

где $\mathbf{s}$ — source_representation (наш полученный честной генерацией $\alpha \cdot \hat{\mathbf{v}}$ ). Один-в-один то же самое, что делает наш ручной hook — но через декларативный API.

Стоп. Зачем тогда нам библиотека — мы все это делали.

Ценность pyvene проявляется в более сложных экспериментах:

Multi-location interventions: одновременно патчим несколько слоёв / компонентов
Causal tracing: зануляем один путь и смотрим, насколько это влияет на output — так находят «важные» слои
Trainable interventions: можно обучить $\mathbf{s}$ под задачу (это то, что делает pyreft , но о нем в следующий сериях)

Итого по способам сделать steering

Все три реализации — PyTorch hooks, nnsight, pyvene — делают одно и то же: перехватывают выход слоя и добавляют $\alpha \cdot \hat{\mathbf{v}}$ .

Cosine similarity между векторами, полученными разными способами, близка к 1 — результат воспроизводим (а для hooks и pyvene наш вектор был один).

Разница между библиотеками в уровне абстракции и количеству промежуточных действий для хорошей интервенции:

Hooks дают контроль и требуют понимание: вы сами регистрируете, модифицируете и снимаете хук. Но легко забыть handle.remove() — и хук висит на модели навсегда; легко перепутать dtype или вернуть не тот tuple. Хороши для одноразовых экспериментов и понимания механики.
nnsight убирает boilerplate (это шаблонный-механический код, который нужно писать каждый раз, но который не несёт смысловой нагрузки). С этой библиотекой хук регистрируется и снимается автоматически, код читается декларативно — «сохрани активацию здесь», «замени её на это». Минус — иногда сложно отладить: ошибка возникает не там, где написана (я по туториалу собрала всё).
pyvene более абстрактен: интервенция описывается как конфиг-объект, а не функция. Это открывает путь к обучаемым интервенциям — можно оптимизировать $\mathbf{v}$ под задачу вместо того, чтобы считать разность средних. На этом построены другие методы.

Для простого steering разницы нет — выбирайте то, что удобнее.

На что обратить внимание? И красивая картинка в финале

В процессе работы мы столкнулись с несколькими из типичных failure mode activation steering. Это хороший момент, чтобы зафиксировать их явно.

Prompt leakage — вектор закодировал не концепт, а поверхностные особенности текста. Tolerant-примеры в датасете лексически похожи на hate: те же субъекты, те же темы, просто с отрицанием. Модель активировала hate-направление на обоих классах, и разность средних указала не туда. Logit lens этого не поймал — он показал hate-токены на минус-стороне и выглядел «правильно». Знак вектора пришлось установить эмпирически.

Oversteering — при $\alpha$ выше ~30 модель теряет связность: начинает повторяться или выдавать мусор. Это сигнал, что активации вышли за пределы тренировочного распределения. Рабочая зона — около 20% от нормы активаций на выбранном слое.

Layer mismatch — один и тот же вектор, построенный на слое 6 и слое 9, даёт разный эффект. На каком слое концепт лучше разделён — вопрос, который решается через PCA (следующий блок) или перебором.

Concept entanglement — наш вектор кодирует ось «агрессивная лексика ↔ формальный регистр», а не строго «hate ↔ tolerant». Это значит, что вместе со сдвигом в сторону hate меняется и стиль — модель становится грубее в целом, не только тематически.

Autoregressive fading и distribution shift мы не измеряли явно, но они присутствуют: эффект на первых токенах сильнее, чем в конце генерации, и вектор, хорошо работающий на eval_prompt, может не работать на произвольных пользовательских запросах (для увидения сего эффекта — просто увеличьте число токенов в output).

Часть из этих проблем можно диагностировать до запуска steering — просто посмотрев на геометрию активаций. Если классы hate и tolerant не разделены в пространстве residual stream, вектор между их центроидами не будет нести концепт — он будет шумом. PCA даёт быстрый визуальный ответ на этот вопрос: хорошо ли линейно разделены классы и совпадает ли направление steering vector с осью разделения. Наша картинка такова:

PC1 и PC2 вместе объясняют только 19.3% дисперсии (10.8% + 8.5%). Это очень мало — значит активации 768-мерного пространства GPT-2 не имеют двух доминирующих осей: дисперсия размазана по многим направлениям, и проекция на плоскость теряет большую часть структуры. Классы сильно перекрываются. Есть слабая тенденция: hate (×) смещён правее по PC1, tolerant (●) — левее и вниз. Но линейного разделения нет. Steering vector направлен примерно между центроидами классов, но не вдоль чёткой разделяющей оси.

PCA подтверждает то, что мы наблюдали эмпирически: концепт hate/tolerant не является доминирующим направлением в residual stream GPT-2 на слое 6. Модель организует своё внутреннее пространство по другим осям — синтаксическим, позиционным, частотным — а hate vs tolerant занимает одно из второстепенных направлений. Стиринг работал, но слабо, именно потому что вектор несёт малую долю полной дисперсии.

Это нормальная картина для маленькой базовой модели. В больших моделях с RLHF или instruction-tuning концепты разделяются чище — там и steering, и интерпретируемость работают лучше. GPT-2 — хорошая учебная площадка именно потому, что здесь всё видно.

Таким образом, activation steering — хоть и сильный, но безумно нюансивный метод, но нюансы стоят того — steering даёт интуицию о том, где и как модель хранит информацию. Это фундамент для более точных методов — probing, causal tracing, SAE. А ещё понимание steering помогает почувствовать архитектуру трансформера — кажется легендарную в наше время.

Отсюда, полезно разобраться в нем руками. И мы с вами прошли путь от сырых PyTorch hooks до декларативных библиотек, столкнулись с реальными failure modes и научились их диагностировать.

Спасибо!

Надеюсь, вы провели время с удовольствием. Если вам понравилось, присоединяйтесь к телеграмм-каналу [Just Data Blog](https://t.me/jdata_blog), ставьте лайки и я не уйду пасти овец, а буду писать новые туториалы.

До встречи!

Ноутбук: GoogleCollab
GitHub: RepoNotebockFile

ссылка на оригинал статьи https://habr.com/ru/articles/1047630/