m23-llm 1.1.0

Mathieu Group M23 Weight Initialization for PyTorch Transformers and LLMs

M23-LLM: Алгебраическая инициализация весов (группа Матьё M23) и диффузионное обучение (dLLM) для GPT-2

Проект исследует альтернативные подходы к обучению и инициализации языковых моделей (LLM). Он объединяет алгебраическую инициализацию весов M23-Spectrum (на основе теории спорадических простых групп и динамической изометрии) и диффузионный режим авторегрессионного декодирования (dLLM по мотивам архитектуры GFusion от Сбера) для модели архитектуры GPT-2 (124M параметров).

---

1. Теоретическая база и архитектура идеи

M23-Spectrum: Теория спорадических групп против случайного хаоса

Классические инициализации (Xavier/He) используют случайные распределения. Это приводит к деформации векторного пространства на глубоких слоях: показатель обусловленности матриц весов (Condition Number) может превышать 1000–6000, из-за чего сигнал сжимается в одних направлениях и растягивается в других.

M23-Spectrum решает эту проблему через концепцию динамической изометрии:

• За основу спектра берутся комплексные корни полинома Элки (связанного с спорадической простой группой Матьё $M_{23}$ порядка $10\ 200\ 960$): $$g^4 + g^3 + 9g^2 - 10g + 8 = 0$$
• Комплексные корни проецируются в вещественную плоскость и циклически распределяются на размерность скрытого слоя ($f_{in}$).
• В спектр вносятся детерминированные микро-возмущения с периодами, соответствующими делителям порядка группы $M_{23}$ ($2, 3, 5, 7, 11, 23$).
• Матрица весов конструируется с помощью SVD-разложения: $$W = U \cdot \text{diag}(\sigma) \cdot V^T$$ Где ортогональные базисы $U$ и $V$ генерируются через QR-разложение, а сингулярные значения $\sigma$ задаются спектром M23.
• Результат: Матрица весов имеет показатель обусловленности равный 1.0. Она идеально сохраняет норму сигналов при проходе через слой (изометрия).

Диффузионный режим обучения (dLLM / GFusion)

Вместо классического авторегрессионного (AR) предсказания следующего токена, модель обучается в режиме диффузии:

• На входе часть токенов маскируется специальным токеном маски ([MASK] / EOS).
• Доля маскируемых токенов выбирается случайно для каждого батча: $t \sim U(0.25, 0.85)$.
• Модель обучается параллельно предсказывать исходные значения только замаскированных токенов (CrossEntropy рассчитывается исключительно по маске).
• Это закладывает базу под нетривиальное понимание контекста и двунаправленное заполнение пропусков с потенциалом ускоренной генерации (Tokens-Per-Forward > 1).

---

2. Аппаратный стек и оптимизация VRAM

Эксперименты проводились локально на ПК со следующей конфигурацией:

• GPU: NVIDIA GeForce RTX 4070 Ti SUPER (16 GB VRAM, архитектура Ada Lovelace)
• RAM: 48 GB DDR5
• ПО: Python, PyTorch 2.7.1+cu126 (с оптимизациями CUDA 12.6 для тензорных ядер Ada Lovelace).

Реализованные оптимизации памяти:

Для предотвращения перегрузки памяти (OOM) и комфортной работы на потребительском GPU были внедрены:

1. Стриминг датасета (Streaming Mode): Обучающий датасет Taiga загружается чанками напрямую из Hugging Face. Это позволило избежать кэширования сотен гигабайт текстовых файлов в оперативной памяти и снизило потребление RAM с 46 ГБ до безопасных 12 ГБ.
2. Смешанная точность (bf16): Использование bfloat16 вместо стандартного fp32 снизило потребление VRAM в два раза без потери стабильности градиентов.
3. Gradient Checkpointing: Включение чекпоинтов градиентов уменьшило пиковое использование VRAM при обучении модели GPT-2 (124M) до ~4.5 ГБ, что позволяет параллельно запускать бенчмарки на одной видеокарте.
4. Оптимизированное QR-разложение: Генерация ортогональных матриц в m23_spectrum.py использует усеченный режим (mode='reduced'), что решило проблему OOM при инициализации огромных эмбеддингов размерностью $50257 \times 768$.

---

3. Быстрый старт и запуск обучения

1. Установка зависимостей

pip install -r requirements.txt

2. Сравнение инициализаций (3000 шагов, бенчмарк-модель)

python compare_init.py

3. Полное обучение (M23 + диффузионный режим dLLM)

python train.py --init_mode m23 --training_mode diffusion --max_steps 10000

4. Обучение в авторегрессионном (AR) режиме для сравнения

# Обучение с M23-инициализацией в AR режиме
python train.py --init_mode m23 --training_mode ar --max_steps 5000

# Обучение с дефолтной инициализацией в AR режиме
python train.py --init_mode default --training_mode ar --max_steps 5000

---

4. Архитектура кодовой базы

Проект реализован в модульной структуре:

• m23_spectrum.py: Математическое ядро. Вычисление корней Элки, наложение гармоник группы Матьё и сборка матриц через SVD и QR.
• m23_init.py: Адаптер инициализации. Выполняет рекурсивный обход модели PyTorch. Корректно определяет типы слоёв (nn.Linear и специфичные для Hugging Face трансформеров Conv1D), производя транспонирование весов для Conv1D. Дополнительно реализует автоматическое масштабирование проекций residual-связей (с коэффициентом $1 / \sqrt{2 \cdot N_{layers}}$) для гашения взрывов энергии.
• model.py: Обертка для GPT-2. Интегрирует переключатель режимов инициализации: m23, orthogonal (стандартная ортогональная), xavier, he и default (стандартная GPT-2 normal).
• dataset.py: Стриминговый загрузчик датасета Taiga (подмножество proza) с реализацией диффузионной маски.
• compare_init.py: Скрипт тестирования сходимости на 3000 шагах. Логирует метрики и генерирует график сравнения.

---

5. Результаты экспериментов (3000 шагов обучения)

Сравнение пяти методов инициализации (M23-Spectrum, Standard Orthogonal, Default GPT-2, Xavier Uniform, He Uniform) на мелкой модели (6 слоев, $d_{embd} = 256$) показало следующие результаты:

Стабильность градиентного потока

На дистанции в 3000 шагов обнаружилось ключевое преимущество ортогональных методов:

• Default GPT-2 (Normal 0.02) после 1000-го шага демонстрирует медленный дрейф нормы градиентов вверх (показатель вырос с 0.8 до ~1.4). Это указывает на постепенный разбаланс residual-потока.
• M23-Spectrum и Standard Orthogonal ведут себя абсолютно стабильно: норма градиентов плавно зафиксировалась в коридоре 0.6–0.8 и шла по идеально ровной горизонтальной линии до конца эксперимента.
• Xavier и He показали стабильность градиентов, но имели катастрофические показатели обусловленности матриц на старте.

Обусловленность матриц весов (Condition Numbers)

• У Xavier Uniform показатель обусловленности для матриц $768 \times 768$ составил ~4400.
• У He (Kaiming) Uniform он составил ~6300.
• У M23-Spectrum и Standard Orthogonal он строго равен 1.0 на всех слоях. Матрицы идеально сохраняют углы и масштабы векторов.

Сходимость (Loss)

• На короткой дистанции (3000 шагов) Default GPT-2 показал чуть меньший Loss (3.131), так как случайный нормальный шум на маленьких масштабах модели облегчает подгонку.
• M23-Spectrum и Standard Orthogonal сошлись плотной группой в районе 3.338–3.400, обогнав Xavier Uniform (3.41). Это подтверждает, что M23 сохраняет все сильные свойства классической ортогональной инициализации.

---

6. Ключевые выводы и ценность проекта

1. Гарантированная изометрия: M23-Spectrum обеспечивает идеальный баланс градиентов без взрывов на старте. Это делает метод перспективным для сверхглубоких сетей (24+ слоев), где обычные инициализации склонны к затуханию сигналов.
2. Спектральная гибкость: В отличие от стандартной ортогональной инициализации, где все сингулярные значения жестко равны 1.0, M23 предлагает структурированный спектр собственных значений на базе группы Матьё. Это сохраняет "разнообразие" сигналов и защищает сеть от коллапса репрезентаций на ранних этапах обучения.
3. Воспроизводимость на бытовом железе: Доказана применимость оптимизаций (bf16 + checkpointing + streaming) для проведения сложных исследований LLM прямо на домашнем ПК с видеокартами игрового сегмента (RTX 4070 Ti SUPER).