AutoEncoders and hidden states Analysis

Привет, друзья!

Я продолжаю цикл туториалов, посвященных области explainable AI. Так, уже были разобраны метод Logit Lens на примере ViT, зондирование gpt2, CAM на примере Yolo NAS — всё можно найти по статьям в профиле. В этом же туториале мы разберем идею применения автокодировщиков для анализа и извлечения признаков из скрытых состояний модели.

В туториале, вы:

• Изучите или повторите, как работает извлечение признаков в Visual Transformers;

• Построите и примените автокодировщик для сжатия скрытых представлений, выученных моделью ViT в задаче классификации котиков и собачек;

• Сравните Vit и PCA в данной задаче.

Приступим!

Шаг 1. Извлечение признаков в ViT

Прежде чем начать анализ скрытых представлений, необходимо понять с чем мы будем работать. Например, в работе с текстовым трансформером, которую я разбирала в туториале Probing GPT model извлечение скрытых представлений базировалось на анализе векторных представлений слов. Там форма скрытых состояний для текстов была такой: (batch_size, sequence_length, hidden_dim). То есть для одного входа, например, для предложения из 6 токенов размерности были (1, 6, 1024) — 6 токенов, каждый представлен 1024-мерным вектором.

Таким образом, при работе с текстом, мы работаем с 1D векторами токенов какой-то фиксированной длины. Но что насчет изображений?

Изображения в признаки. Идея патчей

Интуитивный и базовый подход представления изображения в виде, удобном для анализа и векторизации — взять изображение и сгладить его до 1D. Тогда, если изображение у нас 8×8, то при сглаживании получится вектор длины 64.

Но такой подход не оказывается рабочим, когд изображение чуть больше, скажем 64×64. Тогда его размерность в 1D уже 4096 и при этом каждый пиксель нужно будет связать с другим (то есть связей). Нужно что-то придумать.

Собирая задачи, мы хотим:

1. Чтобы каждое изображение отражали несколько векторов;

2. Чтобы каждый из векторов был фиксированной и адекватной (заданной нами) размерности;

Патчи в ViT.

В ViT (Visual Transformers) — одной из значимых архитектур 2021([статья]) решение этих задач реализовано при помощи патчей.

Патчи строятся так:

1. Берем картинку – высота, ширина, число каналов (1)

2. Режем её на патчей. Каждый из патчей — это “квадрат” на изображении, с размерностями , где — сторона квадрата, — число каналов (2)

3. Эти патчи разглаживаем в вектор, выходит N векторов размерностью (3)

4. Массив патчей, размерностью умножается на обучаемый тензор с размерностями (4)

5. В результате получаем N патчей с размерностями (5)

6. Добавляем [cls] токен, аналогичный BERT (картинка и интуитивный ответ тут, с размерностью , итого у нас вход – патчей (6)

7. В завершение каждому патчу добавляем позиционный эмбеддинг с такой-же размерностью (7)

И вот мы сварили признаки, которые далее проходят в слои энкодера.

Сколько векторов на практике отражает изображение и какова их размерность?

Размерность вектора параметров (то есть d в нотации выше) Vit — это фиксированный параметр. Стандартно, вектор имеет длину 768.

Таким образом, если входное изображение состоит из 1 канала и имеет размерности (h, w) = (224, 224), то чтобы посчитать, сколько векторов будет его характеризовать, нужно выбрать:

1. Количество патчей для разбиения,
2. Размер патчей,

Количество и размер выбираем так, что . Тогда если патчей 14, как в примере, то количество векторов для кодирования изображения:

Разбиение изображения осуществляется при помощи свертки с ядром kernel=(p, p) и шагом (stride), таким что stride=(p, p). На выходе наше изображение будет представлено 197 векторами (196 на патчи, плюс cls) размерностью 768!

Конечно, на практике, писать извлечение руками не надо. Для этого есть ImageProcessor. Работа с ним в коде выглядит так:

MODEL_PATH = 'your_path_here'
MODEL_PATH_HF = 'akahana/vit-base-cats-vs-dogs'

# Загружаем предобученный feature extractor и классификационную модель
feature_extractor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k')
model = ViTForImageClassification.from_pretrained(MODEL_PATH_HF)
model.eval()

# Подготавливаем входные данные
inputs = feature_extractor(images=image, return_tensors="pt")

# изображение после взятия признаков

featured_image = inputs['pixel_values'][0].permute(1, 2, 0)

fig, ax = plt.subplots(figsize=(12, 6))
plt.imshow(featured_image)
plt.axis('off');

Шаг 2. Анализ модели и её структуры

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

P.S. Реализовывать само извлечение вручную нам тоже не придется, поскольку скрытые состояния могут быть получены при помощи указании гиперпараметра при инференсе модели. Поэтому свободно и уверенно выделяем время анализу.

model.vit

Зафиксируем то, что видим — основная модель состоит из последовательности энкодеров и классификатора. Увидим, что будут представлять собой hidden states.

# Делаем предсказание
with torch.no_grad():
    outputs = model(**inputs, output_hidden_states=True)

# Получаем спрогнозированный класс
predicted_class_idx = outputs.logits.argmax(-1).item()
labels = model.config.id2label  # Словарь индексов и классов

print(f"Предсказанный класс: {labels[predicted_class_idx]}")
print(f'Количество скрытых состояний: {len(outputs.hidden_states)}')
print(f'Размерность каждого состояния: {outputs.hidden_states[0].shape}')

# OUTPUT:

# Предсказанный класс: cat
# Количество скрытых состояний: 13
# Размерность каждого состояния: torch.Size([1, 197, 768])

Количество скрытых состояний равно 13, где

• Первое (‘outputs.hidden_states[0]’) — выход части эмбеддинга model.vit.embeddings

• Состояния 2-13 — выходы частей энкодера — все до применения к ним самого последнего слоя компоненты model.vit — layernorm

Важно:

В некоторых моделях, при указании output_hidden_states=True возвращается также last_hidden_state. В общем случае, outputs.last_hidden_state != outputs.hidden_states[-1], так как то, что хранится в outputs.last_hidden_state отражает скрытое состояние после layernorm.

Теперь, поняв выход для одной картинки, мы можем прогнать несколько. Скажем 300 котиков и собачек (можно больше, если запускаете локально на хорошем ноутбуке). Для выполнения, вы можете скачать данные с диска [здесь].

# Указываем путь к папке с изображениями
image_folder = "path_too_folder"

# Подготавливаем списки для данных
last_hidden_states_list = []
mid_hidden_states_list = []
labels_list = []


# Проходим по всем файлам в папке
for filename in tqdm(os.listdir(image_folder)[:300]):
    if filename.endswith((".jpg", ".png", ".jpeg")):  # Фильтруем изображения
        image_path = os.path.join(image_folder, filename)
        image = Image.open(image_path).convert("RGB")  # Открываем изображение

        # Обрабатываем изображение с помощью feature extractor
        inputs = feature_extractor(images=image, return_tensors="pt")

        # Получаем скрытые состояния
        with torch.no_grad():
            outputs = model(**inputs, output_hidden_states=True)
            features = model.vit(**inputs).last_hidden_state[:, 1:, :] # представление изображения перед подачей в классификатор
            mid_features = outputs.hidden_states[len(outputs.hidden_states) // 2][:, 1:, :]  # Последний слой (размерность: batch x patches x hidden_dim)

        predicted_class_idx = outputs.logits.argmax(-1).item()
        labels = model.config.id2label  # Словарь индексов и классов
        label = labels[predicted_class_idx]

        # Добавляем скрытые состояния и метку в список
        last_hidden_states_list.append(features.squeeze(0))  # Убираем batch размерность
        mid_hidden_states_list.append(mid_features.squeeze(0))
        labels_list.append(label)

После запуска кода в ноутбуке, у вас получится следующее:

На каждую картинку из 300 — матрица . В таком формате анализировать признаки и как-то разделить их для двух классов не удобно. Как понять, разделила ли модель классы в пространстве?

Для этого мы можем применить автокодировщик!

Шаг 3. Автокодировщик. Краткая справка.

Автокодировщики (автоэнкодеры) — это нейронные сети особой архитектуры, ключевая компонента которых — извлечение признаков (частный случай, он же классический — сжатие) в скрытое пространство. Изначально это скрытое пространство, конечно, не предназначалось для области XAI (да и в 1986 году, когда идея автокодировщиков была описана, проблемы черных ящиков не было), но последние исследования [OpenAI] (и не только) показали потенциал этой архитектуры в области интерпретируемости. Простейший автокодировщик выглядит [так]. Вы видели его в начале туториала.

Реализуем AE в коде в виде двух компонент, как на картинке — энкодера и декодера. Для них напишем простые классы из такой логики:

У нас 196 патчей по 768 признаков. Количество патчей — это наше колиечество признаков, его урезать некорректно, поэтому будем работать с размерностью скрытого представления — 768.

class PatchEncoder(nn.Module):
    def __init__(self, D_latent=2):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv1d(768, 512, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(512, 256, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(256, 128, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(128, 64, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(64, 32, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(32, 16, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(16, 8, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(8, 4, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(4, D_latent, kernel_size=1), # [B, 196, D_latent].

        )

    def forward(self, x):
        x = x.permute(0, 2, 1)  # [B, 768, 196] для Conv1d
        return self.encoder(x).permute(0, 2, 1)  # [B, 196, D_latent]

    # Аналогичный декодер в ноутбуке

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

При применении автокодировщика к скрытым состониям идея будет такова:

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

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

Шаг 4. Подготовка данных.

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

Почему именно так?

Постановка задачи извлечения признаков всегда связана с извлечением информации из модели, а значит корректно опираться только на ту информацию, которая заложена в ней. При анализе признаков гипотеза такова — спрогнозированный класс соответствует “мнению” модели на основе её знаний об изображении. Скрытые состояния и классы мы уже извлекли, так что просто наведем красоту в class labels.

# Подготовка классов

func = lambda x: 1 if x =='cat' else 0
inv_func = lambda x: 'cat' if x == 1 else 'dog'

labels_list = [func(i) for i in labels_list]

Обучим автоэнкодер.

autoencoder = PatchAutoencoder(D_latent=2)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(autoencoder.parameters(), lr=1e-4)
num_epochs = 10

for epoch in range(num_epochs):

    optimizer.zero_grad()
    outputs = autoencoder(last_hidden_states_tensor)  # [B, 196, 768]
    loss = criterion(outputs, last_hidden_states_tensor)
    loss.backward()
    optimizer.step()
    print(loss)

Шаг 5. Извлечение признаков и анализ

Когда у нас всё есть, мы можем, наконец, посмотреть внутрь. Для этого извлечем скрытые представления и возьмём медиану по патчам, чтобы получить матрицу , где m — число строк (примеров), а n — медиана каждого патча. Обратите внимание — мы обучились на представлениях последнего слоя.

latent = autoencoder.encoder(last_hidden_states_tensor)
latent_med = np.median(latent.detach().numpy(), axis=1)
latent_med.shape

plt.figure(figsize=(8,6))
sns.scatterplot(x=latent_med[:, 0], y=latent_med[:, 1], hue=[inv_func(i) for i in labels_tensor.detach().numpy()], alpha=0.5)
plt.title("2D Visualization of Latent Space")
plt.show()

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

Результаты и выводы

Зафиксируем, что получилось.

1. Сначала была взята обученная модель.

2. На тестовых данных для этой модели были извлечены:
   прогнозы;
   скрытые состояния последнего слоя;
   скрытые состояния среднего слоя;

3. На скрытых состояниях был обучен автокодировщик, с центральной размерностью между encoder и decoder компонетами равной 2.

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

Визуализации показали, что модель хорошо разделяет признаковые представления на последнем скрытом состоянии и практически не разделяет на среднем. Из этого можно выдвинуть гипотезу о постепенном выделении при уточнении глубины сети (причем средний слой достаточно плохо разделяет классы).

Примечание:

На самом деле, вывод этого частного исследования нельзя обобщать для всех моделей и предполагать, что серединные hidden states не кодируют богатые представления. Напротив, например, тут [Layer by Layer: Uncovering Hidden Representations in Language Models] авторы показывают, что промежуточные слои могут кодировать даже более богатые представления, чем последние, показывая хорошую производительность в широком диапазоне подзадач. Ещё интересный пример, здесь [Probing Latent Subspaces in LLM for AI Security: Identifying and Manipulating Adversarial States] извлекают направление, дающее джейлбрейк, также с промежуточного слоя.

Внутри много интересного =).

Заметка на полях – PCA

Давайте рассмотрим, действительно ли нам нужен был автоэнкодер.

Метод главных компонент (PCA) – еще один эффективный метод снижения размерности, который можно было бы применить в этом контексте, особенно если мы пытаемся упростить представление скрытых состояний. PCA – это линейный метод, который проецирует данные в пространство, где дисперсия данных максимальна вдоль новых осей, что делает его хорошим кандидатом для определения наиболее информативных компонентов в пространстве признаков.

Хорошей альтернативой автоэнкодеру PCA делают:

1. Простота: PCA прост и эффективен с точки зрения вычислений. В отличие от автоэнкодеров, которые требуют обучения, PCA можно применять непосредственно к матрице признаков с минимальной подготовкой.

2. Интерпретируемость: PCA предоставляет компоненты, которые напрямую соответствуют направлениям максимальной дисперсии в пространстве признаков.

3. Уменьшение размерности: Как и автоэнкодер, PCA может уменьшать размерность пространства признаков. Разница в том, что PCA делает это линейно, тогда как автоэнкодер способен улавливать нелинейные связи в данных. Однако линейные методы, такие как PCA, могут по-прежнему работать достаточно хорошо, если данные линейно разделимы.

4. Время и ресурсы: PCA не требует обучения, что экономит время и вычислительные ресурсы по сравнению с обучением автоэнкодера, особенно если набор данных большой.

В случае скрытых состояний из модели трансформатора (трансформера) применение PCA может помочь уменьшить сложность данных и потенциально выявить четкие разделения между классами. Проецируя данные в пространство с меньшей размерностью (тоже два или три компонента), мы можем построить график данных и визуально проверить разделимость между классами, как мы пытались сделать с автоэнкодером.

original_hid_states = np.median(last_hidden_states_tensor.detach().numpy(), axis=1)

pca = PCA(n_components=2, random_state=42)
orig_latent_2d = pca.fit_transform(original_hid_states)

# Plot
plt.figure(figsize=(8,6))
sns.scatterplot(x=orig_latent_2d[:, 0], y=orig_latent_2d[:, 1], hue=[inv_func(i) for i in labels_tensor.detach().numpy()], alpha=0.5)
plt.legend()
plt.title("2D Visualization of Latent Space")
plt.show()

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

Результат показывает, что важно помнить традиционные методы, особенно при работе с более простыми моделями и задачами. В них отношения могут быть относительно простыми, так что классические методы могут сэкономить время и вычислительные ресурсы, предоставляя прозрачное и интерпретируемое представление данных. Они ещё служат отличной базой, которая помогают проверить, действительно ли необходимы более сложные модели или достаточно ли более простых методов для эффективного решения проблемы.

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

На этом сегодня всё!

Как всегда, туториал на гитхаб (там же ноутбуки и другие туториалы в репозитории). Надеюсь, в скором времени оформлю похожее по SAE.

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