От учителя к студенту: путь к изотропным эмбеддингам в 29M модели. Часть 0-10

Часть 0. Введение

Важное предупреждение:
Весь код в последующих частях по моей методологии написан нейросетью DeepSeek, сам бы не осилил (я не программист). За мной лишь отладка, исправление неточностей и ошибок.
Тексты этих статей собраны из “заметок на полях” во время реализации, потому возможны некоторые шероховатости в аргументации и стилистике.
Если вас это устроит, желаю приятного чтения.

0.1. Проблема и постановка задачи

Учитель — DeepPavlov/rubert-base-cased-sentence: 180M параметров, 12 слоёв, 768-dim. Качественная модель, но тяжёлая для CPU-инференса и не беспроблемная: у неё есть negation blindness (плохо различает “устроился / не устроился”) и анизотропия (все эмбеддинги сжаты в узкий конус — любые два текста дают высокий cos). То же и у базы sergeyzh/rubert-tiny-turbo, только сильнее: 312-dim с mean pooling сжимает ещё плотнее.

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

Применения: дедубликация текстов, семантический поиск, кластеризация.

Ограничения:

• Инференс на CPU (без GPU).

• Размер модели — до 300 MB.

• Быстродействие — в разы быстрее учителя.

0.2. Четыре цели проекта

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

1. Исследовательская — понять на практике, что даёт внимание ^[1] в моделях. Архитектура эмбеддингов для этого удобнее всего: модель маленькая, своих слоёв внимания нет (только self-attention энкодера, который мы не трогаем), а вся новизна — в pooling-голове поверх готовых hidden states. Если attention pooling даёт заметный эффект — его видно сразу.

2. Дистилляция представлений. Проверить, можно ли за счёт дистилляции от 12-слойного 768-dim учителя получить студента с сопоставимой семантикой. Цель — не просто догнать учителя по метрике retrieval на конкретном датасете, а добиться высокого косинусного сходства (val_cos) с учителем на отложенной выборке.

3. Negation blindness и анизотропия. И turbo, и учитель плохо различают сложные негативы и страдают от анизотропии. Хотелось дать студенту лучшую поддержку различения близких текстов и развести пространство.

4. Практический инструмент. Лёгкая модель для CPU, которая решает три задачи: поиск, дедубликация, кластеризация. Это разные аспекты семантики: поиск — общая семантика, дедубликация — калибровка абсолютных cos, кластеризация — крупнозернистая структура. Модель, которая справляется со всеми тремя, — универсал.

0.3. Метрики успеха

• Cosine к учителю после Phase A — ≥0.90; финальный — ≥0.85.

• Retrieval accuracy на сложных негативных примерах (hard negatives) (hnp_val) — ≥0.97 (turbo со старта 0.968, планка изначально была 0.85, но по факту достигнута до обучения ^[2]).

• Дедубликация: порог, при котором число пар для автоудаления — единицы (не десятки тысяч).

• Кластеризация: отношение intra/inter — заметно выше 1.

Финальные результаты и сравнение с baseline — в разделе 10.

0.4. Архитектурные решения

Базовая модель — sergeyzh/rubert-tiny-turbo. Мы не берём чистый cointegrated/rubert-tiny2, хотя поначалу предполагалась именно эта модель. Причина: turbo — contrastive-pretrained версия tiny2, это даёт более высокий старт по семантике при той же структуре и компактности. Подробное сравнение — в разделе 1.

Дополнительные компоненты:

• Multi-head Attention Pooling (4 головы, выходы голов усредняются) — вместо mean pooling. Обучаемые query, softmax-взвешивание токенов.

• Проекция Linear(312 → 768), bias=True — для дистилляции (сравнение с учителем в его пространстве). После обучения остаётся как опция: projected=True возвращает 768-dim, projected=False — 312-dim (меньше размер, быстрее).

• Токенизатор — родной (84K vocab). Не расширяем.

• Whitening — в основной модели не применяем: contrastive-обучение само даёт близкое к изотропному пространство. Опционально реализуем fit_domain() — whitening-адаптацию под домен пользователя, на случай доменного сдвига.

Что в итоге: 29.4M параметров, 312-dim, 117.8 MB (state_dict), ~120 MB полная папка с токенизатором и обвязкой.

0.5. План по частям

• Окружение и разведка моделей — venv, inspect_layers.py: сравнение трёх моделей по gap на семантических парах, выбор pooling-стратегии.

• Данные: Вики и домен — parse_wiki.py, psycho_corpus.py, build_corpus.py: сборка общего корпуса (200K wiki + 50K psycho = 243,412 текстов после дедубликации) и domain_indices.npy для Phase B.

• Датасет contrastive — prepare_hnp.py: deepvk/ru-HNP (99,999 троек после предобработки: query / 5 pos / 5 neg).

• Создание ученика — build_student.py, check_student_init.py: StudentTiny (turbo + attention pooling + проекция), sanity-проверки.

• Базовая линия — eval_models_hnp.py, eval_cos_to_teacher.py: retrieval на hnp_val для учителя, turbo и студента, cos к учителю на val.txt до обучения.

• Кэширование — precompute_teacher.py, tokenize_corpus.py, tokenize_hnp.py: эмбеддинги учителя и токенизация всех датасетов для Colab.

• 7.1. Phase A — дистилляция. Замороженный энкодер, обучаем только голову (cosine loss к учителю).

• 7.2. Phase B — гибрид. Разморозка верхних 2 слоёв, contrastive (InfoNCE) + distillation (cosine) 1:1, композитный score для чекпоинта.

• 8.1. Внешняя валидация — eval_external.py: ru-paraphrase-NMT-Leipzig (выборка 5000 пар), Top1/Top5/MRR.

• 8.2. Практические задачи — eval_tasks.py, eval_tasks_baseline.py: поиск, дедубликация, кластеризация — на студенте, turbo и учителе.

• Сборка — export_hf.py: самодостаточная папка rubert_tiny2_embed/ с классом SemanticEmbedder (encode, search, dedup, cluster, fit_domain).

• Итоги — что получилось, чему научились, что дальше.

Все прогоны сохраняют параметры и метрики в meta/*.json — по ходу текста мы на них ссылаемся.