meta-attention is all you need

• Важные замечания
  
• Знакомимся с мета-трансформерами
  
• Подробный разбор компонентов
  
• Подробный разбор обучения
  
• Эксперименты
  
• Архитектурные схемы
  
• Заключение
  
• Исходники

• Хуки активации - механизм считывания активации. Срабатывает автоматически когда forward проходит через "свой" слой, извлекает нужную позицию скрытых состояний и складывает в буфер активаций.
  
• Когнитивный энкодер - небольшая нейросеть, превращающая активации из буфера в когнитивные токены. Две основные архитектуры - это линейные проекторы со слоя на токен + небольшая MLP голова, и мини-трансформер. Обе сети показывали эффективные результаты но в разных аспектах, об этом будет рассказано позднее.
  
• Врата внимания (Gates) - обучаемый скалярный множитель, по одному на слой. Регулирует насколько сильно мета-внимание подмешивается в слой — то есть нужна ли ему интроспекция в принципе.
  
• Головы мета-внимания - позволяет отдельно взятому слою выборочно выбирать, к активациям каких других слоёв "прислушиваться" сильнее, а каких слабее. Т.е. он может смотреть на слой A больше чем на слой B.

• Проход 1 — forward без генерации. Хуки активации снимают активации со всех слоёв. Энкодер проецирует их в когнитивные токены и кладёт в буфер.
  
• Проход 2 — forward с активной мета-инъекцией. На каждом слое мета-внимание видит когнитивные токены из буфера и подмешивает мета-сигнал в основной поток через врата. Модель генерирует ответ.

def hook(module, input, output):
    if self._frozen:
        return
    hidden_states = output[0] if isinstance(output, tuple) else output
    # [batch, seq_len, hidden_dim] → берём последний токен
    last_token = hidden_states[:, -1, :].detach().clone()
    self.activations[f"layer_{layer_idx}"] = last_token.squeeze(0)

• Хук срабатывает автоматически когда forward проходит через "свой" слой
  
• Получает полный тензор скрытых состояний [batch, seq_len, hidden_dim]
  
• Извлекает срез последнего токена [:, -1, :] — для авторегрессивной модели это точка принятия решения, hidden state по которому предсказывается следующий токен
  
• .detach() — отвязывает от графа базовой модели (мы не хотим градиенты в базу), .clone() — копия чтобы не держать ссылку на буфер
  
• Складывает в словарь по индексу слоя

# Per-layer проектор (один на каждый из 32 слоёв):
nn.Sequential(
    nn.LayerNorm(hidden_dim),          # 4096
    nn.Linear(hidden_dim, bottleneck), # 4096 → 256
    nn.GELU(),
)
# Encoder gate (скаляр на слой):
nn.Parameter(torch.tensor([0.3]))      # tanh-gated
# Общий output projector:
nn.Sequential(
    nn.LayerNorm(bottleneck),          # 256
    nn.Linear(bottleneck, hidden_dim), # 256 → 4096
    nn.GELU(),
    nn.Linear(hidden_dim, hidden_dim), # 4096 → 4096
)

активация слоя i [4096]
  → projector_i (LayerNorm + Linear → 256 + GELU)
  → encoder_gate_i:  proj * tanh(gate_i)
  → стек по всем 32 слоям → [batch, 32, 256]
  → output_proj (256 → 4096 → GELU → 4096)
  → output_norm (LayerNorm)
  → когнитивные токены [batch, 32, 4096]

self.gate = nn.Parameter(torch.tensor([gate_init], dtype=torch.float32))  # init 0.3
# ...
gate_value = torch.tanh(self.gate)
return residual + gate_value * cross_attention_output

• tanh'(0.3) = 0.91 — почти линейная зона, градиенты текут свободно
  
• tanh'(2.0) = 0.07 — gates замерзают навсегда (мёртвый градиент)
  
• init=0.1 в bfloat16: precision ~0.01, мелкие обновления теряются

Слой      tanh(gate)   Роль
L0-L5     0.01-0.04    токенизация — мета-сигнал не нужен
L6-L12    0.04-0.05    синтаксис, низкоуровневая семантика
L19-L24   0.07-0.09    рассуждение — начинает слушать
L25-L28   0.07-0.11    формирование ответа — активно использует
L29       0.12-0.19    ЛИДЕР — точка принятия решения
L30-L31   0.07-0.08    финальная обработка

# Проекции (bottleneck_dim=256, num_heads=4, head_dim=64):
self.norm     = nn.LayerNorm(hidden_dim)          # pre-norm
self.down_proj = nn.Linear(hidden_dim, bottleneck, bias=False)  # 4096→256
self.q_proj   = nn.Linear(bottleneck, bottleneck, bias=False)   # из сжатого hidden
self.k_proj   = nn.Linear(hidden_dim, bottleneck, bias=False)   # из cog tokens
self.v_proj   = nn.Linear(hidden_dim, bottleneck, bias=False)   # из cog tokens
self.up_proj  = nn.Linear(bottleneck, hidden_dim, bias=False)   # 256→4096
self.token_preference = nn.Parameter(torch.zeros(num_cognitive_tokens))

residual = hidden_states
h = LayerNorm(hidden_states)
h_compressed = down_proj(h)              # [batch, seq, 256]
Q = q_proj(h_compressed)                  # из текущего скрытого состояния
K = k_proj(cognitive_tokens)              # из когнитивных токенов
V = v_proj(cognitive_tokens)
# multi-head: разбить на 4 головы по 64
scores = Q · Kᵀ / √64                      # [batch, heads, seq, 32]
scores = scores + token_preference         # learnable bias на источники
attn   = softmax(scores)
out    = attn · V                          # взвешенная сумма cog токенов
out    = up_proj(out)                      # обратно в 4096
output = residual + tanh(gate) · out       # ← врата здесь

• Прогоняем базовую (замороженную) модель forward'ом на промпте
  
• Хуки снимают активации последнего токена со всех 32 слоёв → [32, 4096]
  
• Записываем активации + правильный ответ + флаг pass1_correct (угадала ли модель сама, без рефлексии)

активация слоя i [4096]
  → LayerNorm + Linear(4096 → 256)
  → ConfidenceHead (256 → 1)
  → P(ответ правильный)

Pass 1 (read):
  base_model.forward(prompt)         # хуки активны, генерации нет
  activations ← хуки [32 × 4096]
  cognitive_tokens ← encoder(activations)   # [32, 4096]
  buffer.fill(cognitive_tokens)
Pass 2 (write + loss):
  хуки замораживаются (freeze)
  logits ← base_model.forward(prompt + target,
                              cross_attention=active)  # головы видят буфер
  loss = CrossEntropy(logits, target_text)
Backward:
  loss.backward()                    # через frozen base → CA → cog tokens → encoder
  optimizer.step()                   # обновляет ТОЛЬКО обвязку

if pass1_correct:
    # CONFIRM: модель сама угадала → подтверждаем
    target = " B) 4 Hz"
    action = "confirm"
else:
    if random() < 0.5:
        # CORRECT: модель ошиблась → учим исправлять
        target = " Wait, the correct answer is B) 4 Hz."
        action = "correct"
    else:
        # REFUSE: модель ошиблась → учим отказываться
        target = " I'm not confident enough to answer this question accurately."
        action = "refuse"

base model:        Llama-3.1-8B-Instruct (bf16, frozen)
learning rate:     2e-4 (×5 для gates/preferences)
batch size:        2, grad accumulation 16 → effective 32
epochs:            10 (early stop patience 5)
max_seq_len:       256
scheduler:         cosine, warmup 5%
dataset:           full MMLU, 12042 train / 1000 val / 1000 test
correction ratio:  0.5
init:              из Phase 1 Selective checkpoint (warm start)
trainable params:  ~188M (encoder 51.7M + 32 CA 136.5M)
frozen:            8.0B base

• Frozen base обязателен. Любая разморозка базы (LoRA, partial unfreeze) даёт shortcut: модель оптимизирует loss напрямую через свои веса, минуя мета-канал. Refusal rate коллапсирует с 9.2% до 0.4%. Опровергнуто 10 экспериментами на Gemma-2B.
  
• Gate init в линейной зоне tanh. init=0.3 → tanh'(0.3)=0.91 (градиент течёт). init=2.0 → tanh'(2.0)=0.07 (gates замерзают навсегда). Критичная деталь которая определяет учатся gates вообще или нет.
  
• Pretrain проекторов = обязательное предварительное условие для глубоких энкодеров. Без него 32-слойная архитектура не сходится.
  
• Сложность задачи как гиперпараметр. На лёгких задачах (TriviaQA, baseline 76%) gates закрываются до 0.01 — канал не нужен. На сложных (MMLU Hard, baseline 40%) gates стабилизируются на 0.08-0.12. Модель адаптивно регулирует использование интроспекции в зависимости от того, нужна ли она.
  
• Быстрая сходимость. 2 эпохи до лучшего результата. Обучаем проводку, не знания — поэтому быстро.