Что такое logit lens (интерпретация скрытых состояний)?

Краткий тезис

Logit lens — это техника интерпретации внутренних представлений трансформеров, которая проецирует скрытые состояния (hidden states) промежуточных слоёв на выходное словарное пространство (vocabulary) через unembedding матрицу. Это позволяет увидеть, какое слово (или токен) модель «думает» на каждом слое, не дожидаясь финального слоя. Logit lens даёт интуитивное понимание того, как модель постепенно формирует ответ: от неопределённых эмбеддингов на ранних слоях к уверенному предсказанию на поздних.

1. Зачем нужна интерпретация скрытых состояний

Современные LLM — это «чёрные ящики»: мы видим вход и выход, но не знаем, что происходит внутри. Для отладки, улучшения и доверия к модели важно понимать, как она обрабатывает информацию. Скрытые состояния (hidden states) — это векторные представления на каждом слое, которые несут информацию о контексте. Logit lens — один из самых простых и наглядных методов «заглянуть внутрь».

Основные цели

• Понять, на каком слое модель «решает», каким будет следующее слово.
• Выявить нежелательные паттерны (например, предвзятость на ранних слоях).
• Сравнить поведение разных моделей или слоёв.

2. Как работает logit lens: базовая идея

В трансформере каждый слой выдаёт state|скрытое состояние ( h^{(l)} ) (вектор размерности ( d_{[text](/wiki/text){model}} )). Финальный слой (( L )) преобразуется в логиты (logits) через unembedding матрицу ( W_U ) (размер ( |V| \times d_{[text](/wiki/text){model}} )), а затем в вероятности через softmax:

[ [text](/wiki/text){logits}^{(L)} = h^{(L)} \cdot W_U^T, \quad p = [text](/wiki/text){softmax}([text](/wiki/text){logits}^{(L)}) ]

Logit lens предлагает применить ту же unembedding матрицу к скрытому состоянию любого промежуточного слоя ( l ):

[ [text](/wiki/text){logits}^{(l)} = h^{(l)} \cdot W_U^T ]

Полученные логиты интерпретируются как «предсказание» модели на слое ( l ). Если на слое ( l ) модель уже «знает» ответ, то argmax логитов будет совпадать с финальным предсказанием.

Ключевой момент unembedding матрица — это линейное преобразование, которое «переводит» state|скрытое состояние в пространство токенов. Она не меняется от слоя к слою, поэтому проекция даёт единую шкалу для сравнения.

3. Математическая основа и детали

Unembedding матрица ( W_U ) — это часть выходного проекционного слоя (обычно последний линейный слой перед softmax). В большинстве трансформеров (например, GPT-2, Llama) она является транспонированной версией embedding матрицы ( W_E ) (weight tying), но не всегда.

Формула для одного токена на позиции ( t ):

[ [text](/wiki/text){logits}^{(l)}_t = h^{(l)}_t \cdot W_U^T ]

где ( h^{(l)}_t ) — скрытое состояние на слое ( l ) для позиции ( t ).

Проблема скрытые состояния на разных слоях имеют разную норму и смещение (bias). Logit lens игнорирует layer norm и residual stream — на практике это может приводить к тому, что на ранних слоях argmax даёт мусорные токены. Тем не менее, для средних и поздних слоёв метод часто работает удивительно хорошо.

Пример кода (PyTorch-like):

def logit_lens(model, hidden_states, layer_idx):
    # hidden_states: (batch, seq_len, d_model) для слоя layer_idx
    W_U = model.lm_head.weight  # unembedding matrix
    logits = hidden_states @ W_U.T
    return logits

4. Что показывают слои на примере Llama

Для моделей семейства Llama (и других decoder-only) наблюдается характерная картина:

Важно точные номера слоёв зависят от модели. Для GPT-2 (12 слоёв) картина аналогична, но сжата.

Пример визуализации для запроса «Столица Франции —» на слое 10 модель может показывать «Париж» с вероятностью 0.3, а на слое 12 — уже 0.9.

5. Сравнение с другими методами интерпретации

Logit lens — это «бесплатный» метод, не требующий дополнительного обучения. Он хорошо подходит для быстрой диагностики.

6. Ограничения и альтернативы

Основные ограничения

• Unembedding bias скрытые состояния на ранних слоях могут иметь другое распределение, чем на последнем, поэтому проекция даёт смещённые логиты.
• Layer norm: в современных моделях (например, GPT-2) перед unembedding применяется final layer norm. Logit lens без учёта этого даёт менее точные результаты.
• Residual stream: скрытое состояние — это сумма вкладов всех предыдущих слоёв. Logit lens не разделяет эти вклады.

Альтернативы

• Tuned lens (Nostalgebraist, 2020): обучает аффинное преобразование (сдвиг + масштаб) для каждого слоя, чтобы лучше соответствовать финальному распределению.
• Softmax lens: применяет softmax к проекции, но с учётом температуры.
• Logit lens с layer norm: перед проекцией применяет final layer norm к скрытому состоянию (если архитектура позволяет).

Когда logit lens работает плохо

• Для очень глубоких моделей (70B+), где внутренние представления сильно отличаются.
• Для задач, требующих многошаговых рассуждений (chain-of-thought) — модель может «думать» на нескольких слоях одновременно.

7. Практическое применение в RAG и агентах

В контексте Agentic RAG logit lens может помочь:

• Отладка генерации понять, на каком слое модель начинает использовать контекст из retrieval.
• Выявление «галлюцинаций» если на средних слоях модель «думает» неправильный токен, а на поздних исправляется — это может указывать на конфликт между контекстом и внутренними знаниями.
• Оптимизация длины контекста если модель «понимает» ответ уже на 10-м слое, можно обрезать вычисления (early exit).

Пример: для запроса «Какой город является столицей Франции?» с контекстом «Париж — столица Франции» logit lens покажет, что уже на 5-м слое модель склоняется к «Париж», а на 8-м — уверена.

8. Связь с residual stream и attention

Residual stream — это сумма скрытых состояний всех предыдущих слоёв. Logit lens применяется к этому суммарному вектору. Это означает, что предсказание на слое ( l ) отражает совокупное знание, накопленное до этого слоя.

Attention влияет на residual stream через взвешенное суммирование. Logit lens не показывает, какие именно токены contributed, но можно комбинировать с attention: например, смотреть, на какие токены модель обращала внимание, когда «решила» предсказать слово.

9. Пет-проект для закрепления

Задача Реализовать logit lens для небольшой модели (например, GPT-2) и визуализировать, как меняется предсказание по слоям для заданного промпта.

Инструменты

• Python, PyTorch, transformers (Hugging Face)
• Matplotlib для визуализации

Шаги:

1. Загрузить предобученную модель GPT-2 (gpt2).
2. Выбрать промпт, например: «The capital of France is».
3. Прогнать промпт через модель, сохраняя скрытые состояния всех слоёв (использовать output_hidden_states=True).
4. Для каждого слоя ( l ) применить unembedding матрицу (model.lm_head.weight) к скрытому состоянию последнего токена.
5. Взять softmax и получить вероятности для топ-5 токенов.
6. Построить график: ось X — номер слоя, ось Y — вероятность правильного токена (например, « Paris»).
7. Интерпретировать: на каком слое модель «уверена» в ответе.

Ожидаемый результат График, показывающий рост вероятности правильного токена от случайной (ранние слои) до высокой (поздние слои). Для неправильных промптов (например, «The capital of France is London») можно увидеть, что модель «колеблется» между токенами.

Дополнительно Реализовать tuned lens — обучить аффинное преобразование для каждого слоя на небольшом корпусе и сравнить с обычным logit lens.

10. Связь с другими вопросами

11. Навигация

• Предыдущий: 673
• Следующий: 675
• Индекс: 00. Индекс разборов

Навигация

• Предыдущий: 673
• Следующий: 675
• Индекс: 00. Индекс разборов