Привет, друзья! Этот туториал посвящён зондированию (probing) — простому, но мощному методу для изучения внутренней работы LLM (больших языковых моделей). С его помощью можно получить приближенные знания о паттернах, которые выучивает модель и о том, как эти знания распространяются по слоям.

Метод простой, но довольно интересный. К туториалу прилагается ноутбук с кодом. Всех заинтересовавшихся — прошу к чтению!

Идея зондирования

Техники «зондирования модели» (в литературе — probing) — это множество простых идейно и эффективных методов анализа внутренних представлений. Смысл процесса зондирования исходит из определения слова — «зондирование» (английский вариант «probing») — это изучение, исследование или изучение чего‑либо глубоким способом — на основе погружения. Так, можно зондировать почву — лезть в глубину грунта или желудок — лезть в глубину желудка.

Идея напрашивается сама собой — давайте залезем внутрь модели!

Как мы туда полезем?

Подадим в модель входной текст , получив прогноз . Сохраним и рассмотрим цепочку скрытых состояний:

$x \to hidden_1 \to hidden_2 \to \dots \to hidden_n \to output$

где — векторное представление входного из на слое.

Если применить к , собранным с определенных данных модели, можно сформировать гипотезы для ответа на вопросы:

Есть ли в семантическая информация о частях речи? (собираем представления и решаем задачу классификации на аннотированных данных)
Где (на каких уровнях в модели) закодированы знания о фактах или концептах, которые модель «выучила» из данных? (аналогично предыдущему)

Справедливо заметить: зондирование требует рутинной разметки, однако она не всегда сложна.

В туториале рассмотрим:

Процесс зондирования на примере GPT2;
Анализ информативности скрытых состояний с помощью PCA;
Постановку эксперимента (и сам эксперимент) для ответа на вопрос: какой слой по уровню позволяет приближенно решить задачу регресси и хранит информацию по годам?

Постановка задачи — где в модели хранятся знания о датах?

Пусть у нас есть генеративная модель. Поставим вопрос: где именно эта модель «знает», в какие годы жил тот или иной человек?

Для этого нам нужно:

Создать набор данных: пары вида (вопрос: «Когда родился Ньютон?», ответ: «1643»).
Попустить вопросы через модель и извлекаем скрытые состояния с разных слоёв.
Извлечь спрогнозированную дату;
Обучить зонд, который предсказывает дату на основе скрытого состояния.
Проанализировать, на каких слоях модель наиболее эффективно хранит информацию о датах.

Приступим.

Модель GPT-2

Использовать будем GPT-2. GPT-2 — это трансформер, состоящий из блоков (transformer layers). Для начала загрузим модель и посмотрим на пример её работы. В ходе анализа будем использовать gpt2-medium, обученный на WebText. Подробно про модель можно почитать в [статье], но для себя зафиксируем, что особенности датасета, на котором обучена модель, таковы:

Нет страниц википедии;
Знания для модели собраны до 2019 года;
Сбор данных проводился на основе статей на Reddit с фильтрацией по пользовательскому голосу >3;

Особенности модели всегда нужно учитывать при постановке эксперимента.

# Загружаем токенизатор и модель
tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
model = GPT2LMHeadModel.from_pretrained('gpt2-medium')

# Установка устройства (GPU, если доступно)
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model = model.to(device)

# Исходный текст
text = "When was Albert Einstein Born?"
encoded_input = tokenizer(text, return_tensors='pt') #.to(device)

# Генерируем продолжение текста
with torch.no_grad():
  # = model(**encoded_input)
  output_ids = model.generate(**encoded_input, max_length=20)

# Декодируем токены обратно в текст
decoded_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)

print(decoded_output)

Также, зафиксируем, что когда мы передаём текст в модель, происхожит следующее:

Текст разбивается на токены (например, "Albert" → 15433, "Einstein" → 8372).
Каждый токен превращается в векторное представление (embedding).
Вектора передаются через несколько слоёв трансформера. Последние слои добавляют всё больше информации о контексте.
На выходе каждого слоя получаем последовательность скрытых состояни1 и итоговый прогноз, которые постепенно «насыщаются» контекстом.

Подробный процесс этого приложен кодом в ноутбуке. Также зафиксируем, что модель с генерацией ответа, выолняет, в нашем случае, задачу ответа на вопрос. Постепенно доработаем код так, чтобы из ответа извлекать только год. Ограничимся годами с 1800 до 2019.

Извлекать год будем с помощью регулярных выражений.

pattern = r"\b(1[89][0-9]{2}|20[01][0-9])(?:[s’']{0,2}|th)?\b"
example = "In the 2000s, I fell in love with cats."

match_ = re.search(pattern, example)

year = match_.group(1) if match_ else None

print("Extracted year:", year)

Подготовка данных и сбор hidden states

Скрытые состояния (hidden states) — это внутренние представления входного текста внутри модели. Они формируются после обработки данных каждым слоем нейросети и содержат закодированную информацию о тексте.

«Что содержат скрытые состояния?» — открытый вопрос. Есть ряд исследований, показывающий что на разных уровнях содержится разная семантическая информация. Например, на ранних слоях — информация о частях речи и месте слова в предложении, а на поздних — значение слова. Но эта интерпретация не универсальна и в частности зондирование — попытка понимания внутренних процессов в модели.

Как получить скрытые состояния в Hugging Face?

Если модель поддерживает output_hidden_states=True, то после обработки текста мы можем достать их так:

outputs = model(**encoded_input)
hidden_states = outputs.hidden_states  # Это кортеж из N тензоров (по числу слоёв)

hidden_states[0] — первый слой
hidden_states[-1] — последний слой
hidden_states[len(hidden_states) // 2] — центральный слой

Форма скрытых состояний: (batch_size, sequence_length, hidden_dim). Для gpt2-medium: (1, 6, 1024) — 6 токенов, каждый представлен 1024-мерным вектором.

# Исходный текст
texts = ["When was Albert Einstein born?", "When was Frida Kahlo born?", "When was Claus Hammel born?"]
encoded_input = tokenizer(texts, return_tensors='pt', padding=True, truncation=True)

# Прогоняем через модель и получаем скрытые состояния
with torch.no_grad():
    outputs = model(**encoded_input, output_hidden_states=True)
    output_ids = model.generate(**encoded_input, max_length=30)
    hidden_states = outputs.hidden_states  # Все скрытые состояния (tuple из 25 слоёв)

Набор данных

Будем работать с [people dataset], предобработанным и очищенным для данного эксперимента. Датасет, необходимый для работы вы можете скачать с [google disc].

most_popular = pd.read_csv('//Users/sabrinasadieh/Code/XAI-open_materials/gpt2_probing/people_dataset_prepared_most_popular.csv')
most_popular.head()

И так, у нас есть набор данных, предварительно очищенный по табличке "год". В нем содержается люди, рожденные с 1800 по 2019. Нас будет интересовать две колонки — 'name' (для извлечения имени) и birth.

Сбор скрытых состояний

Как это работает?

Выбор данных: Мы выбираем входные данные и метки для вспомогательной задачи. Например, это может быть задача определения частей речи, извлечения временных рамок или классификации фактов.
Извлечение скрытых состояний: Пропускаем вход через модель и сохраняем скрытые состояния с одного или нескольких слоёв.
Обучение зонда: строим простой классификатор или регрессионную модель (например, логистическую регрессию, SVM или небольшой нейросетевой слой), обучая её на . Эта модель и есть наш «зонд».
Анализ: Оцениваем производительность зонда. Если зонд решает задачу хорошо, это значит, что информация, необходимая для решения задачи, закодирована в .

# Функция для получения активаций и внимания
def get_activations_and_attention(model, enc_inputs, pattern_to_response):
    activations = {}

    with torch.no_grad():
        outputs = model(**enc_inputs, output_hidden_states=True)
        output_ids = model.generate(**enc_inputs, max_length=30)

    decoded_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    match_ = re.search(pattern_to_response, decoded_output)
    year = match_.group(1) if match_ else None

    # Извлекаем внимание (первый, средний и последний слои)
    activations['layer_1'] = outputs.hidden_states[0]  # Первый слой внимания
    activations['layer_middle'] = outputs.hidden_states[len(outputs.hidden_states) // 2]  # Средний слой внимания
    activations['layer_last'] = outputs.hidden_states[-1]  # Последний слой внимания

    return activations, year

Анализ активаций

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

X_acivations_l1 = pd.read_csv('results/X_acivations_l1.csv')
X_acivations_lmid = pd.read_csv('results/X_acivations_lmid.csv')
X_acivations_llast = pd.read_csv('results/X_acivations_llast.csv')

В данных содержатся вопросы из 6-ти токенов "[SOS] When was {name} born? [EOS]". Почитстим данные от дублей.

print(X_acivations_l1.duplicated().sum()) # 504
print(X_acivations_lmid.duplicated().sum()) # 504
print(X_acivations_llast.duplicated().sum()) # 504

X_acivations_l1.drop_duplicates(inplace=True)
X_acivations_lmid.drop_duplicates(inplace=True)
X_acivations_llast.drop_duplicates(inplace=True)

Всего у нас 1024*6 + 2 столбца. Каждые 1024 — есть кодирование токенов. Проанализируем дубли по столбцам, соотвествующим одному токену в последовательности. Получаем следующее:

Слой 1.

Токен 1, активаций-дублей: 4996, процент: 99.98
Токен 2, активаций-дублей: 4996, процент: 99.98
Токен 3, активаций-дублей: 3952, процент: 79.09
Токен 4, активаций-дублей: 3345, процент: 66.94
Токен 5, активаций-дублей: 3263, процент: 65.3
Токен 6, активаций-дублей: 3570, процент: 71.44

Центральный слой.

Токен 1, активаций-дублей: 4996, процент: 99.98
Токен 2, активаций-дублей: 4994, процент: 99.94
Токен 3, активаций-дублей: 3612, процент: 72.28
Токен 4, активаций-дублей: 1544, процент: 30.9
Токен 5, активаций-дублей: 601, процент: 12.03
Токен 6, активаций-дублей: 248, процент: 4.96

Последний слой.

То же, что и на центральном.

Заметим, что:

Активации первых токенов ([SOS] When) в 99% случаев одинаковы.
Активации токена 3 ('was') одинаковы в более чем 70% случаев
"Перелом" уникальности — токен 4 — начало имени персоны.

Везде выкинем первый-второй. Остальные — выкидывать не будем, так как нам также важна комбинация токенов.

Обучение зондирующих моделей.

Обратим внимание на размерность пространства признаков. Заметим, что модель сейчас обучать не совсем корректно — высока вероятность переобучения (а если бы признаков было больше, чем объектов — обучать было бы совсем не корректно).

Размерность сократим при помощи метода главных компонент, так, чтобы доля объясняемой дисперсии была 0.7, 0.8, 0.9, 0.95 соответственно. Посмотрим, сколько компонент для этого нужно.

По стандарту — выборку делим на train-test. Для корректного оценивания важно, чтобы на каждом слое были выбраны одинаковые индексы при разбиении на train/test. Этого можно достичь, зафиксировав random_state в train_test_split из sklearn.

# Функция для применения PCA и вывода дисперсии
def get_pca_data(X_train, X_test, n_components=100):
    pca = PCA(n_components=n_components)
   
    X_train_pca = pca.fit_transform(X_train)
    X_test_pca = pca.transform(X_test)

    print(f'Varince explained: {np.sum(pca.explained_variance_ratio_)} with ({n_components} components)')

    return X_train_pca, X_test_pca

Чтобы достичь разного уровня объясненной дисперсии на первом слое, необходимо компонент на слое 1:

для объясненной дисперсии 0.7 = 480
для объясненной дисперсии 0.8 = 700
для объясненной дисперсии 0.9 = 1100
для объясненной дисперсии 0.95 = 1500

На центральном слое:

для объясненной дисперсии 0.7 = 175
для объясненной дисперсии 0.8 = 310
для объясненной дисперсии 0.9 = 600
для объясненной дисперсии 0.95 = 930

На последнем:

для объясненной дисперсии 0.7 = 2
для объясненной дисперсии 0.8 = 4
для объясненной дисперсии 0.9 = 7
для объясненной дисперсии 0.95 = 17

Чем глубже слой, тем более значимо меньше нам нужно компонент. Причем, скорее всего, где-то в модели есть переломный момент — когда уменьшение количества необходимых компонент перестает значимо снижаться.

Обучение зонда

Обучим модели на уровне сохраненной дисперсии 0.9. Прогнозировать будем года, спрогнозированные моделью, а не истинные, поскольку модель не всегда дает верный ответ и, кроме того, если его не дает, значит бесполезно искать знания о нём внутри.

def get_probing_results(model, X_train, X_test, y_train, y_test, params='unknown'):

    model.fit(X_train, y_train)
    predictions = model.predict(X_test)

    mae = mean_absolute_error(predictions, y_test)
    mse = mean_squared_error(predictions, y_test)
    r2 = r2_score(y_test, predictions)

    metrics = [mae, mse, r2]

    print(f'Metrics of experiment with params {params}, n_components {X_test.shape[1]} are: \n')
    print(f'MAE: {metrics[0]}')
    print(f'MSE: {metrics[1]}')
    print(f'R2: {metrics[2]}')

    return model, predictions, metrics

Лучшая модель по результату получается на среднем слое. Это позволяет выдвинуть гипотезу о том, что, если в модели отдельно представлена информации, связанныя с числами, то она представлена ближе к центру. Запуская код в ноутбуке, вы заметите, что $R^2$ стабильно меньше или около нуля. Посмотрим на прогнозы, чтобы проанализировать это.

Прогнозы модели истинной модели кусочны, большая часть сконцентрирована около 1860-1900 годов. Прогнозы же линейных моделей похожи на нормальное распределение.

Эти факты делают второе слагаемое в большим, так как в целом линейная модель требует нормального распределения целевой переменной.

Можно попробовать decision tree. Оно переобучается, но дает картинку лучше.

Выводы

На этом туториал завершен. Зафиксируем главное: анализ зонда позволяет приблизиться к оценке расположения знаний внутри модели.

В случае эксперимента в туториале, лучшими оказались средние слои, что не всегда справедливо.

Также хочу заметить, что сам анализ зондов можно улучшать, например строить не линейную модель (например, дерево) или ограничить анализ прогнозами до 1900 года (так как других прогнозов всего 8.9% от данных).

Надеюсь, туториал станет хорошей отправной точкой для вас!

Благодарю за ваше время!
Туториал на гитхаб (там же ноутбуки).

Успешных вам проектов!
Ваш Дата-автор!

Probing GPT model