Grafted Titans: как заставить малую модель запоминать на 90% больше контекста

Когда контекста не хватает, а видеопамять не резиновая

Представьте: у вас есть Qwen-2.5-0.5B — компактная модель, которая в теории должна работать на скромном железе. Но вот беда — её контекстное окно ограничено 8192 токенами. Вы пытаетесь обработать длинный документ, и модель просто... забывает начало. Знакомо?

Традиционные решения? Либо увеличивать контекст (что требует тонны VRAM), либо использовать RAG (что добавляет сложность). А что если сделать модель умнее? Заставить её эффективнее использовать ту память, что уже есть?

Grafted Titans: не расширяем окно, а учимся в нём жить

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

1 Как работает этот «привитый титан»

Архитектура состоит из трёх ключевых компонентов:

• Graft Module — отдельный трансформерный блок, который параллельно основной модели обрабатывает контекст
• Cross-Attention Gating — механизм, который решает, какую информацию взять из основного потока, а какую — из graft
• Test-Time Training — адаптация параметров graft'а под конкретную задачу во время инференса

class GraftedTitans(nn.Module):
    def __init__(self, base_model, graft_dim=256):
        super().__init__()
        self.base_model = base_model
        self.graft = TransformerBlock(dim=graft_dim)
        self.gate = nn.Linear(base_model.hidden_dim + graft_dim, 2)
        
    def forward(self, x, context):
        # Основной поток
        base_out = self.base_model(x)
        
        # Параллельный graft поток
        graft_out = self.graft(context)
        
        # Динамическое смешивание
        gate_weights = F.softmax(self.gate(torch.cat([base_out, graft_out], dim=-1)), dim=-1)
        
        return gate_weights[:, :, 0:1] * base_out + gate_weights[:, :, 1:2] * graft_out

Цифры, которые заставляют задуматься

Авторы протестировали Grafted Titans на BABILong benchmark — наборе задач, которые специально проверяют способность модели работать с длинными контекстами.

Девяносто один процент улучшения на том же контекстном окне. Не увеличивая размер модели в разы, не требуя 24 ГБ видеопамяти. Просто добавив умный адаптер.

Test-Time Training: обучение на лету

Самая интересная часть — как graft адаптируется. Вместо того чтобы заранее обучать его на миллионах примеров, он учится прямо во время работы.

2 Как это выглядит в коде

def train_graft_during_inference(model, input_sequence, target_output, steps=10):
    """
    Адаптируем graft под конкретную задачу прямо во время инференса
    """
    # Замораживаем основную модель
    for param in model.base_model.parameters():
        param.requires_grad = False
    
    # Размораживаем только graft
    for param in model.graft.parameters():
        param.requires_grad = True
    
    optimizer = torch.optim.Adam(model.graft.parameters(), lr=1e-4)
    
    for step in range(steps):
        optimizer.zero_grad()
        output = model(input_sequence)
        loss = F.cross_entropy(output, target_output)
        loss.backward()
        optimizer.step()
        
        if step % 5 == 0:
            print(f"Step {step}: loss = {loss.item():.4f}")
    
    return model

Звучит безумно? Обновлять параметры модели во время инференса? Но это работает. Graft учится выделять паттерны из конкретного контекста, который вы ему дали.

Сравнение с альтернативами: почему не RAG и не просто большая модель?

Давайте честно: RAG — это костыль. Элегантный, работающий, но всё же костыль. Вам нужно:

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

Grafted Titans решает проблему на уровне архитектуры модели. Нет внешних систем, нет дополнительных компонентов. Просто модель, которая стала лучше помнить.

А что насчёт просто взять модель побольше? Если вас интересует масштабирование на несколько карт, то да, можно взять Qwen-2.5-7B или 14B. Но тогда вы теряете возможность запуска на скромном железе.

Практическое применение: где это реально нужно

Представьте, что вы строите локального AI-ассистента на GTX 1650. У вас 4 ГБ VRAM. Варианты:

1. Использовать крошечную модель с контекстом 4K — она будет быстро забывать
2. Использовать модель среднего размера с RAG — сложно настраивать
3. Взять Qwen-0.5B с Grafted Titans — и получить почти вдвое лучшую память

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

Как НЕ надо использовать Grafted Titans

Не пытайтесь применить это к уже огромным моделям типа Llama 3.1 405B. Там и так достаточно параметров для хорошей памяти. Идея в другом — сделать маленькие модели умнее, а не большие — ещё больше.

Не используйте Test-Time Training на продакшене без кэширования. Обучение на каждом запросе съест все ваши ресурсы. Лучший подход — адаптировать graft один раз под тип документов, а потом использовать закэшированную версию.

Под капотом: cross-attention gating в деталях

Механизм gate — это то, что отличает Grafted Titans от простого конкатенирования выходов. Вместо того чтобы тупо складывать результаты основного блока и graft'а, модель учится динамически взвешивать их важность.

# Плохой подход (что многие делают)
combined = torch.cat([base_output, graft_output], dim=-1)
output = self.projection(combined)  # Просто линейный слой

# Хороший подход (Grafted Titans)
gate_weights = torch.sigmoid(self.gate_network(base_output, graft_output))
# gate_weights — тензор формы [batch, seq_len, 2]
# где первое измерение — вес для base_output
# второе — для graft_output
output = gate_weights[:, :, 0:1] * base_output + gate_weights[:, :, 1:2] * graft_output

Gate network учится отвечать на вопрос: «Для этого конкретного токена, в этом конкретном контексте — что важнее: информация из основной модели или из graft'а?»

Кому подойдёт эта архитектура

Если вы попадаете в одну из этих категорий — присмотритесь к Grafted Titans:

• Разработчики edge-устройств: у вас мало памяти, но нужно работать с длинными контекстами
• Исследователи: экспериментируете с архитектурами моделей и хотите улучшить маленькие LLM
• Энтузиасты локальных LLM: устали от ограничений контекстного окна на домашнем железе
• Стартапы с ограниченным бюджетом: не можете позволить себе инференс больших моделей, но нуждаетесь в хорошей памяти

Если же вы просто хотите запустить чат-бота и у вас есть RTX 4090 — возможно, проще взять модель побольше. Но если вы как раз выбираете GPU для первого AI-PC с ограниченным бюджетом, то Grafted Titans даст вам больше за те же деньги.

Ограничения и подводные камни

Ничто не идеально. У Grafted Titans есть свои тёмные стороны:

1. Test-Time Training требует времени: первые N запросов будут медленнее, пока graft адаптируется
2. Дополнительные параметры: хоть и немного, но они есть. Если у вас совсем туго с памятью, каждый мегабайт на счету
3. Сложность отладки: теперь у вас две системы внимания, которые могут конфликтовать
4. Не решает все проблемы: если ваша задача требует действительно огромного контекста (100K+ токенов), вам всё равно понадобится что-то вроде RAG или долговременной памяти

Что дальше? Экосистема вокруг нейронной памяти

Grafted Titans — не единственный эксперимент в этой области. Похожие идеи появляются в других проектах:

• Memory Networks: отдельные сети, которые специализируются на запоминании
• Differentiable Neural Computers: гибрид нейросетей и внешней памяти
• Neural Turing Machines: попытка дать нейросетям доступ к «ленте», как у машины Тьюринга

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

Что если применить эту архитектуру к другим маленьким моделям? К Genesis-152M-Instruct? К Phi-2? Результаты могут быть ещё интереснее.

Собираем всё вместе: минимальный рабочий пример

import torch
import torch.nn as nn
import torch.nn.functional as F

class GraftModule(nn.Module):
    """Простой graft модуль на основе трансформера"""
    def __init__(self, dim=256, num_heads=8):
        super().__init__()
        self.attention = nn.MultiheadAttention(dim, num_heads)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        
    def forward(self, x):
        attn_out, _ = self.attention(x, x, x)
        x = self.norm1(x + attn_out)
        ffn_out = self.ffn(x)
        return self.norm2(x + ffn_out)

class GraftedQwen(nn.Module):
    """Qwen-2.5-0.5B с graft модулем"""
    def __init__(self, base_model):
        super().__init__()
        self.base_model = base_model
        hidden_dim = base_model.config.hidden_size
        
        # Graft модуль с меньшей размерностью
        self.graft = GraftModule(dim=256)
        
        # Проекция для согласования размерностей
        self.graft_proj = nn.Linear(256, hidden_dim)
        
        # Gate network
        self.gate = nn.Linear(hidden_dim * 2, 2)
        
    def forward(self, input_ids, attention_mask=None):
        # Основной forward pass
        base_output = self.base_model(input_ids, attention_mask=attention_mask)
        
        # Graft forward pass (параллельно)
        # Упрощённо: используем embeddings как вход для graft
        embeddings = self.base_model.get_input_embeddings()(input_ids)
        graft_output = self.graft(embeddings.transpose(0, 1)).transpose(0, 1)
        graft_output = self.graft_proj(graft_output)
        
        # Динамическое смешивание
        combined = torch.cat([base_output.last_hidden_state, graft_output], dim=-1)
        gate_weights = F.softmax(self.gate(combined), dim=-1)
        
        # Взвешенная сумма
        mixed = (gate_weights[:, :, 0:1] * base_output.last_hidden_state + 
                 gate_weights[:, :, 1:2] * graft_output)
        
        return BaseModelOutput(last_hidden_state=mixed)

# Использование
from transformers import AutoModelForCausalLM

# Загружаем базовую модель
base_qwen = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B")

# Оборачиваем в Grafted Titans
model = GraftedQwen(base_qwen)

# Test-Time Training (упрощённо)
def adapt_to_document(model, document_tokens, steps=5):
    optimizer = torch.optim.Adam(model.graft.parameters(), lr=1e-4)
    
    for step in range(steps):
        optimizer.zero_grad()
        output = model(document_tokens)
        # Здесь должна быть ваша функция потерь
        # Например, loss = language_modeling_loss(output, document_tokens)
        loss.backward()
        optimizer.step()
    
    return model

Это упрощённый пример, но он показывает основную идею. На практике нужно аккуратно обрабатывать attention masks, правильно инициализировать graft и настраивать гиперпараметры обучения.

Стоит ли игра свеч?

Если вы боретесь с ограничениями контекстного окна на малом железе — определённо да. Grafted Titans даёт +91% точности на тестах долгой памяти ценой +0.8 ГБ VRAM. Это хороший trade-off.

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

Самое важное — Grafted Titans показывает, что можно улучшать маленькие модели, не делая их больше. Не добавляя миллиарды параметров. Не требуя экзотического железа. Просто переосмысливая, как они используют то, что уже есть.

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