induction_heads

Концепция "smeared keys" и induction heads из статьи "In-context Learning and Induction Heads"

Цель исследования

Сравнение эффективности моделей Conv-Attn и Attn-Attn в:

1. Аппроксимации "smeared keys"
2. Формировании induction heads
3. Качестве next-token prediction

Методология

Архитектуры моделей

Компонент	Conv-Attn	Attn-Attn
Эмбеддинг	nn.Embedding	nn.Embedding
Основной слой	Conv1d(k=3) + MultiheadAttention	2x MultiheadAttention
Гиперпараметры	lr=1e-3, 4 heads	lr=2e-3, 4 heads

Датасет и обучение

• Датасет: Wikitext-2 (36718 примеров)
• Токенизация: GPT-2 tokenizer (max_length=128)
• Критерий: CrossEntropyLoss
• Эпох: 25
• Батч: 32

Результаты

Метрики обучения

Модель	Start Loss	Final Loss	Start Acc	Final Acc
Conv-Attn	1.3955	0.0070	81.94%	99.81%
Attn-Attn	2.5709	1.2687	67.60%	75.10%

Детальный анализ

Conv-Attn Model

Особенности внимания:

1. Сильная вертикальная структура
2. Доминирование последнего токена
3. Слабые диагональные элементы

Вес свертки:

• Диапазон: [-0.04, 0.04]
• Локальные паттерны, соответствующие smeared keys

1. Локальная свертка (kernel_size=3) → идеально для предсказания next-token в локальном контексте
2. Causal padding → сохраняет позиционную информацию
3. Меньше параметров → стабильнее обучение

Attn-Attn Model

Слой 1:

• Четкие диагонали (до 0.6)
• Классические induction heads
• Паттерны "A→B→A"

Слой 2:

• Пик самовнимания (0.6 в позиции 6-6)
• Хаотичный фон (<0.05)
• Признаки коллапса внимания

Ключевые выводы

1. Эффективность архитектур:

   • Conv-Attn показала лучшие результаты
   • Attn-Attn демонстрирует более теоретически правильные механизмы внимания

2. Паттерны внимания:

   • Conv-Attn: вертикальные структуры + позиционный bias
   • Attn-Attn: Layer1 - диагонали, Layer2 - самовнимание

3. Проблемы и решения:

Проблема	Возможное решение
Коллапс внимания	Добавить masked_fill диагонали
Позиционный bias	Causal padding в свертке
Низкая точность Attn	Оптимизировать lr второго слоя