Латентное пространство вместо токенов: собираем экспериментальную LLM по частям

Зачем ломать то, что работает? Потому что токены — это тупик

Все современные LLM построены на токенах. GPT, Llama, Mistral — все они жуют текст, разбитый на кусочки. Это удобно? Для компьютера — да. Для смысла — катастрофа.

Представьте, что вы читаете книгу, но каждое слово заменено номером из словаря. «Война и мир» превращается в последовательность типа [45, 128, 7, 893...]. Модель видит эти номера и пытается угадать следующий. Она не оперирует смыслами. Она играет в угадайку с цифрами.

Мой эксперимент начался с простого вопроса: а что, если работать не с токенами, а сразу с векторами смысла? Не переводить «кошку» в токен 1234, а сразу в вектор [0.12, -0.45, 0.78...]. Так родилась идея LLM с латентным пространством. Если вам интересна философия этого подхода, я подробно разбирал её в статье Прочь от токенов: как я строил LLM с латентным пространством.

Архитектура: три головы, одна цель

Классическая трансформерная модель — монолит. Моя архитектура — модульный конструктор из трёх независимых блоков. Каждый решает свою задачу, и главное — их можно обучать отдельно.

Магия в разделении труда. Энкодер и декодер можно взять готовыми и заморозить — они уже умеют переводить текст в векторы и обратно. Вся работа по «мышлению» ложится на ядро. И оно в разы меньше классической LLM.

1 Выбор энкодера: почему USER-bge-m3, а не что-то другое

Энкодер — краеугольный камень всей архитектуры. От него зависит, насколько хорошо текст превращается в вектор. Плохой энкодер — и вся модель будет генерировать бессмыслицу.

Я перепробовал десятки вариантов:

• Sentence-BERT: Хорош, но слишком общий. Не понимает инструкций.
• E5: Отличный для поиска, но для генерации — слабоват.
• Текст-энкодеры от OpenAI: Дорого, закрыто, и не факт, что лучше.

Остановился на USER-bge-m3. Почему?

1. Он обучен на инструкциях. Понимает разницу между «Объясни» и «Переведи».
2. Размерность 1024 — золотая середина. Меньше — теряем информацию. Больше — взрываем VRAM.
3. Открытые веса. Можно качать, менять, дообучать.

from sentence_transformers import SentenceTransformer

# Загружаем энкодер - это наша "глазная сетчатка"
encoder = SentenceTransformer('BAAI/bge-m3', trust_remote_code=True)
encoder = encoder.encode  # Для удобства

# Теперь любой текст превращается в вектор 1024D
vector = encoder("Привет, как дела?")
print(f"Вектор размерности: {vector.shape}")  # (1024,)

2 Сердце системы: проектируем латентное ядро

Ядро — это маленькая модель, которая работает исключительно с векторами. На входе — последовательность векторов (например, 10 векторов по 1024 числа). На выходе — один вектор, предсказание следующего.

Архитектура ядра — упрощённый трансформер:

• Линейный слой для проецирования 1024 → 512 (экономия VRAM)
• 4 слоя внимания с 8 головами
• LayerNorm и dropout для стабильности
• Линейный слой 512 → 1024 обратно

import torch
import torch.nn as nn

class LatentCore(nn.Module):
    def __init__(self, input_dim=1024, hidden_dim=512, num_layers=4, num_heads=8):
        super().__init__()
        self.proj_in = nn.Linear(input_dim, hidden_dim)
        self.proj_out = nn.Linear(hidden_dim, input_dim)
        
        # Упрощённый трансформер-энкодер
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim*4,
            dropout=0.1,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
    def forward(self, x):
        # x: [batch, seq_len, 1024]
        x = self.proj_in(x)  # -> [batch, seq_len, 512]
        x = self.transformer(x)
        x = self.proj_out(x)  # -> [batch, seq_len, 1024]
        return x[:, -1, :]  # Берём последний вектор как предсказание

Почему так мало параметров? Потому что ядро не должно быть большим. Его задача — научиться логике последовательностей векторов, а не нужно, оно уже есть в encoder_layer

Почему именно такая архитектура? Потому что она в 50 раз меньше Llama 7B. И при этом способна улавливать зависимости между векторами. Проверено.

3 Декодер: от вектора обратно к словам

Самая сложная часть. Нужно превратить вектор 1024D в осмысленный текст. Здесь два пути:

1. Использовать готовую маленькую LLM как декодер (например, TinyLlama). Проектировать вектор в её эмбеддинг-пространство и запускать генерацию.
2. Обучить свой мини-декодер с нуля. Сложнее, но даёт полный контроль.

Я выбрал второй путь. Почему? Потому что готовые LLM ожидают токены, а у нас вектор. Проектирование вектора в эмбеддинг-пространство — ещё один источник ошибок.

class VectorToTextDecoder(nn.Module):
    """Простой декодер: вектор -> последовательность токенов"""
    def __init__(self, latent_dim=1024, vocab_size=32000, max_seq_len=128):
        super().__init__()
        self.latent_proj = nn.Linear(latent_dim, 512)
        self.token_embeddings = nn.Embedding(vocab_size, 512)
        
        # Небольшой трансформер для генерации
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=512, nhead=8, batch_first=True
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=3)
        
        self.output_layer = nn.Linear(512, vocab_size)
        
    def forward(self, latent_vector, target_tokens=None):
        # latent_vector: [batch, 1024]
        batch_size = latent_vector.size(0)
        
        # Проектируем вектор в пространство эмбеддингов
        memory = self.latent_proj(latent_vector).unsqueeze(1)  # [batch, 1, 512]
        
        if target_tokens is not None:
            # Режим обучения: учимся предсказывать следующий токен
            tgt_emb = self.token_embeddings(target_tokens)
            output = self.transformer(tgt_emb, memory)
            return self.output_layer(output)
        else:
            # Режим генерации (упрощённо)
            # Здесь нужен авторегрессионный цикл
            pass

Обучение по частям: главный трюк для экономии VRAM

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

Фаза 1: Замораживаем энкодер

USER-bge-m3 уже отличный. Его не трогаем. Просто используем как чёрный ящик, который превращает наш датасет в векторы.

# Подготовка датасета: текст -> векторы
from datasets import load_dataset

ds = load_dataset("your_dataset")

def encode_batch(batch):
    # Энкодер работает в батчах для скорости
    vectors = encoder(batch["text"])
    return {"vectors": vectors}

ds_encoded = ds.map(encode_batch, batched=True, batch_size=32)
ds_encoded.save_to_disk("./dataset_encoded")

Фаза 2: Обучаем ядро на векторах

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

Loss функция: Cosine similarity loss или MSE. Я использовал комбинацию:

def latent_loss(pred_vector, target_vector):
    # 1. Косинусная близость для направления
    cos_loss = 1 - F.cosine_similarity(pred_vector, target_vector).mean()
    
    # 2. MSE для абсолютных значений
    mse_loss = F.mse_loss(pred_vector, target_vector)
    
    return cos_loss + 0.1 * mse_loss  # Веса подбираются экспериментально

На этой фазе потребление VRAM — смешное. Ядро на 10 миллионов параметров спокойно обучается на RTX 4090 с 24 ГБ. Для сравнения: Llama 7B требует минимум 48 ГБ для обучения. Если вы хотите копнуть глубже в железные вопросы, посмотрите мою статью про сервер для LLM на Radeon.

Фаза 3: Обучаем декодер

Самый ресурсоёмкий этап, но всё равно в разы легче, чем обучать полную LLM. Берём пары (вектор → текст) и учим декодер превращать векторы обратно в слова.

Здесь есть тонкость: нужно использовать teacher forcing и маскирование, как в обычных языковых моделях.

Генерация текста: как это работает на практике

Когда все три модуля обучены, генерация выглядит так:

1. Пользователь вводит текст: «Расскажи анекдот про программистов»
2. Энкодер превращает его в вектор V1
3. Ядро принимает [V1] и генерирует следующий вектор V2 (предполагаемый ответ)
4. Декодер превращает V2 в текст: «Программист звонит в техподдержку...»

Для более длинных ответов можно итеративно подавать сгенерированный текст обратно в энкодер и продолжать цепочку.

Где собака зарыта: главные проблемы и как их обойти

Экспериментальная архитектура — это всегда грабли. Вот на какие я наступил:

А оно вообще работает? Результаты эксперимента

После двух месяцев проб и ошибок модель заработала. Не так хорошо, как GPT-4, конечно. Но:

• Отвечает на вопросы осмысленно (в пределах простых тем)
• Поддерживает диалог на 5-10 реплик
• Занимает на диске 800 МБ вместо 14 ГБ (Llama 7B)
• Требует для инференса 4 ГБ VRAM вместо 12+ ГБ

Главное достижение — доказательство концепции. Латентное пространство работает. Модульная архитектура работает. Раздельное обучение экономит ресурсы.

Что дальше? Куда развивать архитектуру

Если вы решите повторить эксперимент или пойти дальше, вот направления для улучшений:

1. Диффузионное ядро. Вместо авторегрессивного трансформера использовать диффузионную модель для генерации векторов. Это может дать более разнообразные и креативные ответы. Сравнение архитектур я делал здесь.
2. Мультимодальность. Энкодеры бывают не только для текста. Добавить изображения, аудио — и получится единое латентное пространство для всего.
3. Квантование с обучением. Использовать VQ-VAE для превращения непрерывных векторов в дискретные коды. Это упростит ядро (оно станет похоже на классическую LLM, но с другим «словарём»).
4. Специализированные энкодеры. Обучить энкодер для конкретной области (медицина, юриспруденция, код). Как в статье про fine-tune под язык программирования, но для энкодера.

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

P.S. Если вы дочитали до этого места и думаете «зачем всё это, если есть ChatGPT» — вы правы. Для 95% задач достаточно готовых моделей. Но эти 5% — там, где нужно что-то уникальное, дешёвое, работающее на странном железе или делающее то, чего большие модели не умеют — вот там такие эксперименты обретают смысл. Как тот самый чемодан без ручки, который тащить тяжело, но бросить жалко.