Как Llama 3.2 3B думает внутри: удивительные находки трассировки активаций

Проблема чёрного ящика: почему мы не понимаем, как думает ИИ?

Современные языковые модели, такие как Llama 3.2 3B, демонстрируют впечатляющие способности, но остаются для нас чёрными ящиками. Мы видим результат — связный текст, ответы на вопросы, решение задач — но не понимаем, как именно модель приходит к этим выводам. Это не просто академический интерес: без понимания внутренних механизмов мы не можем гарантировать безопасность, предсказуемость и надежность ИИ-систем.

Что такое трассировка активаций и как она работает?

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

Когда Llama 3.2 3B обрабатывает запрос \"Сколько будет 2 + 2?\", в её 32 слоях и 3072 нейронах происходят миллионы вычислений. Трассировка активаций позволяет:

• Зафиксировать значения активаций в ключевых точках модели
• Проанализировать паттерны активаций для разных типов запросов
• Выявить специализированные нейронные цепи (circuits), отвечающие за конкретные функции
• Сравнить внутренние представления разных моделей

Этот подход тесно связан с circuit analysis — методом, который ищет минимальные подграфы нейронной сети, необходимые для выполнения конкретной задачи. Если вам интересно сравнение внутренних представлений разных моделей, рекомендую статью \"Как «мыслят» Llama-3 и Qwen-2.5: сравниваем геометрию внутренних представлений\".

Подготовка инструментов для исследования

Для трассировки активаций Llama 3.2 3B нам понадобятся:

Для локального запуска моделей в формате GGUF можно использовать LM Studio или llama.cpp. Важно понимать, что выбор формата модели влияет на точность анализа: например, квантование может искажать активации.

Пошаговый план трассировки активаций

1Загрузка модели и установка хуков

Первым делом загружаем Llama 3.2 3B и регистрируем хуки (hooks) для захвата активаций. Хуки — это функции, которые PyTorch вызывает при прохождении тензора через определённый слой.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np

# Загружаем модель и токенизатор
model_name = \"meta-llama/Llama-3.2-3B\"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)

# Словарь для хранения активаций
activations = {}

# Функция-хук, которая сохраняет активации
def get_activation(name):
    def hook(module, input, output):
        # Сохраняем выходные активации слоя
        activations[name] = output.detach().cpu().numpy()
    return hook

# Регистрируем хуки для всех слоёв декодера
for i, layer in enumerate(model.model.layers):
    layer.register_forward_hook(get_activation(f\"layer_{i}\"))

2Генерация и захват активаций

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

def trace_activations(prompt):
    # Очищаем предыдущие активации
    activations.clear()
    
    # Токенизация и подготовка входа
    inputs = tokenizer(prompt, return_tensors=\"pt\")
    input_ids = inputs[\"input_ids\"]
    
    # Прямой проход с захватом активаций
    with torch.no_grad():
        outputs = model(input_ids)
    
    # Возвращаем словарь активаций и сгенерированный текст
    generated_text = tokenizer.decode(outputs.logits.argmax(-1)[0])
    return activations, generated_text

# Примеры промптов для анализа
prompts = [
    \"Сколько будет 2 + 2?\",
    \"Переведи на английский: Привет, как дела?\",
    \"Объясни, что такое гравитация\",
    \"Напиши хайку про искусственный интеллект\"
]

# Собираем активации для всех промптов
all_activations = {}
for prompt in prompts:
    act, text = trace_activations(prompt)
    all_activations[prompt] = act
    print(f\"Промпт: {prompt[:30]}... | Сгенерировано: {text[:50]}...\")

3Анализ паттернов активаций

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

def analyze_activation_patterns(all_activations):
    # Преобразуем активации в удобный формат
    patterns = {}
    
    # Для каждого слоя вычисляем среднюю активацию по промптам
    for prompt, acts in all_activations.items():
        for layer_name, activation in acts.items():
            if layer_name not in patterns:
                patterns[layer_name] = {}
            
            # Средняя активация по всем нейронам слоя
            patterns[layer_name][prompt] = np.mean(activation)
    
    return patterns

patterns = analyze_activation_patterns(all_activations)

# Визуализация тепловой карты
import matplotlib.pyplot as plt
import seaborn as sns

# Подготовка данных для визуализации
layer_names = list(patterns.keys())
prompt_list = prompts
data = np.array([[patterns[layer][prompt] for prompt in prompt_list] for layer in layer_names])

plt.figure(figsize=(10, 8))
sns.heatmap(data, xticklabels=prompt_list, yticklabels=layer_names, cmap=\"viridis\")
plt.title(\"Паттерны активаций по слоям и промптам\")
plt.xlabel(\"Промпты\")
plt.ylabel(\"Слои\")
plt.tight_layout()
plt.show()

4Выявление нейронных цепей (circuits)

Самый интересный этап — поиск цепей нейронов, которые работают вместе для выполнения конкретной задачи. Мы используем корреляционный анализ между активациями нейронов.

def find_circuits(activations_dict, threshold=0.8):
    circuits = {}
    
    # Анализируем только математический промпт для примера
    math_prompt = \"Сколько будет 2 + 2?\"
    math_activations = activations_dict[math_prompt]
    
    # Для каждого слоя анализируем корреляции между нейронами
    for layer_name, activation in math_activations.items():
        # activation имеет форму [batch_size, seq_len, hidden_size]
        # Усредняем по последовательности
        avg_activation = np.mean(activation, axis=1)[0]  # [hidden_size]
        
        # Ищем сильно коррелирующие нейроны
        # (в реальном анализе нужно анализировать попарные корреляции)
        highly_active = np.where(avg_activation > np.percentile(avg_activation, 90))[0]
        
        if len(highly_active) > 0:
            circuits[layer_name] = highly_active.tolist()
    
    return circuits

math_circuits = find_circuits(all_activations)
print(f\"Найдено {len(math_circuits)} слоёв с высокоактивными нейронами на математические задачи\")
for layer, neurons in list(math_circuits.items())[:5]:
    print(f\"  {layer}: {len(neurons)} активных нейронов\")

Удивительные находки в Llama 3.2 3B

Анализ активаций Llama 3.2 3B раскрыл несколько интересных паттернов:

Нюансы и возможные ошибки

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

1. Корреляция ≠ причинность: Нейроны могут активироваться вместе, но не быть причинно-связанными.
2. Зависимость от промптов: Недостаточно разнообразный набор промптов может создать ложные паттерны.
3. Влияние квантования: Как показано в статье про F16 vs Q8_0, квантованные модели могут иметь искажённые активации.
4. Статистическая значимость: Нужно проверять находки на статистически значимой выборке промптов.
5. Вычислительная сложность: Полная трассировка всех активаций для больших моделей требует значительных ресурсов.

Заключение и дальнейшие шаги

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

Для дальнейшего изучения рекомендую:

• Экспериментировать с разными моделями, включая PLaMo 3 (31B) и Llama 3.3 8B-Instruct
• Изучить более продвинутые методы интерпретируемости, такие как патч-трансплантация и дистилляция цепей
• Применить полученные знания для отладки и улучшения собственных моделей

Понимание внутренних механизмов ИИ — это не только научный интерес, но и практическая необходимость для создания безопасных и надёжных систем. Трассировка активаций — ваш ключ к этому пониманию.