Как работает эмбеддинг слой и почему его размер (embedding dimension) важен?

Краткий тезис

Эмбеддинг слой — это первый обучаемый слой в трансформере, который преобразует дискретные токены в плотные векторные представления фиксированной размерности d (embedding dimension). Размерность d напрямую определяет ёмкость модели (capacity), объём памяти и вычислительные затраты. Выбор d — это компромисс между качеством представления и эффективностью, часто эвристически связанный с размером словаря V (например, d ≈ 1.5 × √V). В контексте Agentic RAG эмбеддинг слой используется как для кодирования запросов и документов, так и для внутреннего представления токенов в LLM-агенте.

1. Что такое эмбеддинг слой и зачем он нужен

Эмбеддинг слой (embedding layer) — это таблица поиска (lookup table) размером V × d, где V — размер словаря (vocabulary size), а d — размерность эмбеддинга. Каждому токену из словаря сопоставляется уникальный плотный вектор (эмбеддинг) размерности d. В отличие от one-hot кодирования (размерность V), эмбеддинги позволяют модели компактно представлять семантику токенов и обучать их в процессе тренировки.

Зачем он нужен

• Преобразует дискретные входные данные (токены) в непрерывное пространство, где можно применять операции линейной алгебры и attention.
• Позволяет модели улавливать семантические и синтаксические сходства между токенами (близкие по смыслу токены имеют похожие векторы).
• Является основой для всех последующих слоёв трансформера.

2. Как работает эмбеддинг слой: архитектура и вычисления

2.1 Таблица эмбеддингов

Эмбеддинг слой хранит матрицу E размером V × d. Каждая строка — это вектор для соответствующего токена. Индексация: для токена с id = i из входной последовательности извлекается строка E[i].

2.2 Прямой проход (forward pass)

1. Вход: последовательность токенов (например, [12, 45, 789]).
2. Для каждого токена выполняется lookup: из матрицы E берётся строка по индексу.
3. Результат: тензор формы (batch_size, seq_len, d).

2.3 Обучение

Матрица E инициализируется случайно (например, из нормального распределения) и обновляется во время обратного распространения ошибки. Градиенты по эмбеддингам считаются стандартно через chain rule. Таким образом, эмбеддинги «учатся» представлять токены в контексте задачи.

2.4 Связь с выходным слоем

Часто в трансформерах веса эмбеддинг слоя и выходного (lm_head) слоя разделяют (tied embeddings), что уменьшает число параметров и улучшает обобщение. В этом случае матрица E используется и для получения логитов на выходе.

3. Размерность эмбеддинга d: что это и как выбирается

Размерность эмбеддинга (embedding dimension) — длина вектора каждого токена. Типичные значения: 512 (BERT-base), 768 (BERT-large), 4096 (Llama-3-70B), 8192 (GPT-4).

3.1 Эвристика «1.5 × √V»

Эмпирическое правило: d ≈ 1.5 × √V, где V — размер словаря. Для V = 128k (как в Llama-3) √128000 ≈ 357.8, 1.5 × 357.8 ≈ 537. Однако на практике d берут гораздо больше (4096), так как современные модели требуют высокой ёмкости.

3.2 Факторы, влияющие на выбор d

• Размер словаря V: чем больше V, тем больше нужно d, чтобы различать токены.
• Глубина модели: в глубоких моделях эмбеддинги проходят через много слоёв, поэтому можно использовать меньшую d (информация «накапливается»).
• Вычислительные ограничения: d определяет размер матриц Q, K, V в attention (d_model = d), а также размер FFN (обычно 4×d). Увеличение d в 2 раза ведёт к ~4-кратному росту FLOPs в attention.
• Качество downstream-задач: для сложных задач (перевод, генерация) требуется большая d.

4. Влияние d на ёмкость модели (capacity)

Ёмкость модели (model capacity) — способность запоминать и обобщать закономерности. Большая d позволяет:

• Хранить более детальные семантические признаки каждого токена.
• Улавливать тонкие различия между похожими токенами.
• Поддерживать большее количество «смысловых измерений» для композиции в attention.

Однако чрезмерно большая d при фиксированном объёме данных ведёт к переобучению (overfitting) и росту числа параметров без улучшения качества.

5. Влияние d на память и вычислительные затраты

5.1 Параметры эмбеддинг слоя

Количество параметров = V × d. Для V=128k и d=4096: 128000 × 4096 ≈ 524 млн параметров — это значительная часть модели (например, в Llama-3-70B общее число параметров ~70B, эмбеддинги ~0.5B).

5.2 Память при инференсе

• Хранение матрицы E: 524M × 2 байта (float16) ≈ 1 ГБ.
• При batch-обработке нужно хранить эмбеддинги для всей последовательности: batch_size × seq_len × d × 2 байта.

5.3 Вычислительные затраты

Lookup — это просто чтение из памяти, но умножения на эмбеддинги не происходит. Однако d влияет на все последующие слои: размерность d_model = d, attention вычисляет QK^T размером seq_len × seq_len с d-мерными векторами, что даёт O(seq_len² × d) FLOPs. FFN имеет размер 4×d, что даёт O(seq_len × d²) FLOPs.

6. Связь с размером словаря и количеством слоёв

6.1 V и d: баланс

Если V большой, а d маленький, эмбеддинги не смогут различить все токены (коллизии). Если V маленький, а d большой — избыточность. Современные LLM используют V ~ 32k–128k и d ~ 4096–8192.

6.2 d и глубина

В моделях с большим числом слоёв (например, 80 слоёв в Llama-3-70B) можно использовать меньшую d, так как каждый слой добавляет нелинейные преобразования. Однако на практике d остаётся высокой для сохранения expressiveness.

7. Примеры из современных LLM

Видно, что с ростом размера модели d растёт, но нелинейно.

8. Особенности эмбеддинг слоя в контексте Agentic RAG

В Agentic RAG эмбеддинг слой используется не только внутри LLM, но и для:

• Кодирования запросов и документов: часто применяют отдельные эмбеддинг-модели (sentence-transformers), но в некоторых архитектурах (например, REPLUG) эмбеддинги берутся из LLM.
• Представления контекста: агент может использовать эмбеддинги токенов для принятия решений (какой документ выбрать, какой инструмент вызвать).
• Fine-tuning: при дообучении RAG-агента эмбеддинг слой обновляется, что улучшает качество retrieval и генерации.

Размерность d эмбеддинг слоя LLM влияет на качество retrieval, если эмбеддинги используются напрямую: большая d даёт более богатые представления, но увеличивает latency и стоимость хранения векторных индексов.

9. Практические рекомендации по выбору d

• Для небольших моделей (до 1B): d = 512–1024.
• Для средних (1B–10B): d = 2048–4096.
• Для больших (10B+): d = 4096–8192.
• Эвристика: d ≈ (0.5–2) × √V, но лучше ориентироваться на бенчмарки и доступные ресурсы.
• Tied embeddings: всегда использовать, если V × d ≈ d_model × V (экономия параметров).

10. Типичные ошибки и заблуждения

• «Больше d — всегда лучше»: нет, из-за переобучения и вычислительных затрат.
• «Эмбеддинги можно не обучать»: в современных LLM они обучаются с нуля, предобученные эмбеддинги (word2vec) устарели.
• «Размерность d не влияет на качество retrieval»: влияет, если эмбеддинги используются для поиска; большая d требует больше памяти для индекса.

Пет-проект для закрепления

Задача: Исследовать влияние размерности эмбеддинга d на качество и производительность небольшой языковой модели.

Инструменты: PyTorch, Hugging Face Transformers, небольшой датасет (например, WikiText-2).

Шаги:

1. Реализуйте простой трансформер с возможностью менять d (d_model).
2. Зафиксируйте V = 10000, число слоёв = 6, heads = 8.
3. Обучите модели с d = 128, 256, 512, 1024 на задаче language modeling (предсказание следующего токена).
4. Замерьте:
   • Perplexity на валидации.
   • Время обучения одной эпохи.
   • Пиковое использование GPU памяти.
5. Постройте графики зависимости perplexity и времени от d.
6. Сделайте вывод: при каком d достигается наилучший trade-off.

Ожидаемый результат: Вы увидите, что после определённого d (например, 512) perplexity перестаёт улучшаться, а время и память растут квадратично.

Связь с другими вопросами

Навигация

• Предыдущий: 670
• Следующий: 672
• Индекс: 00. Индекс разборов