Metadata-Version: 2.4
Name: m23-llm
Version: 1.1.0
Summary: Mathieu Group M23 Weight Initialization for PyTorch Transformers and LLMs
Author-email: M23-Spectrum Team <contact@m23spectrum.dev>
License: MIT
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.19.0
Requires-Dist: torch>=1.9.0
Requires-Dist: transformers>=4.0.0

# 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. Установка зависимостей
```bash
pip install -r requirements.txt
```

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

### 3. Полное обучение (M23 + диффузионный режим dLLM)
```bash
python train.py --init_mode m23 --training_mode diffusion --max_steps 10000
```

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

![Сравнение методов инициализации](comparison.png)

---

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

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