Metadata-Version: 2.4
Name: m23-llm
Version: 1.2.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
License-File: LICENSE
Requires-Dist: numpy>=1.19.0
Requires-Dist: torch>=1.9.0
Requires-Dist: transformers>=4.0.0
Dynamic: license-file

# M23-LLM: Algebraic Weight Initialization & Diffusion Training

<details>
<summary><b>🇬🇧 English (Click to expand/collapse)</b></summary>

# M23-LLM: Algebraic Weight Initialization (Mathieu Group M23) and Diffusion Training (dLLM) for GPT-2, Qwen and LLaMA

The project explores alternative approaches to training and initializing Large Language Models (LLMs). It combines **M23-Spectrum algebraic weight initialization** (based on sporadic simple Mathieu groups and dynamical isometry principles) with a **diffusion autoregressive decoding training mode (dLLM)**.

**Compatibility:** The initialization module is fully adapted for both classical architectures (GPT-2 with Conv1D/Linear layers) and modern LLMs based on SwiGLU activations and RoPE (including **Qwen 3, Qwen 3.5-3.6**, and **LLaMA 3/3.1** models).

---

## 1. Theoretical Concept & Architecture

### M23-Spectrum: Mathieu Group Theory vs. Random Chaos
Classical initializations (Xavier/He) use random continuous distributions. On deep layers, this leads to accumulation of spatial distortions: the condition number of weight matrices can exceed **1000–6000**, compressing the signal in some directions and stretching it in others.

**M23-Spectrum** solves this via **dynamical isometry**:
* The spectrum is derived from the complex roots of the **Elki polynomial** (associated with the sporadic simple Mathieu group $M_{23}$ of order $10,200,960$):
  $$g^4 + g^3 + 9g^2 - 10g + 8 = 0$$
* Complex roots are projected onto the real axis and cyclically mapped to the dimension of the hidden layer ($f_{in}$).
* High-frequency deterministic micro-perturbations are added with periods corresponding to the divisors of the Mathieu group order ($2, 3, 5, 7, 11, 23$).
* The weight matrix is constructed using Singular Value Decomposition:
  $$W = U \cdot \text{diag}(\sigma) \cdot V^T$$
  where orthogonal bases $U$ and $V$ are generated via QR decomposition, and singular values $\sigma$ are defined by the M23 spectrum.
* **Result**: The weight matrix has a condition number **strictly equal to 1.0**. It preserves the norm of signals as they pass through the layers (isometry).

### Diffusion Training Mode (dLLM / GFusion)
Instead of classical autoregressive next-token prediction, the model is trained in a diffusion-like mode:
* Part of the input tokens are masked with a special token (`[MASK]` or `EOS`).
* The ratio of masked tokens is selected randomly for each batch: $t \sim U(0.25, 0.85)$.
* The model is trained to predict the original values of the masked tokens in parallel (CrossEntropy is calculated only on the mask).
* This provides a solid foundation for bidirectional context understanding and non-trivial text completion.

---

## 2. Hardware Stack and VRAM Optimizations

Experiments were conducted on the following configuration:
* **GPU**: NVIDIA GeForce RTX 4070 Ti SUPER (16 GB VRAM, Ada Lovelace)
* **RAM**: 48 GB DDR5
* **Software**: Python, PyTorch 2.7.1+cu126 (CUDA 12.6 optimizations for tensor cores).

### Implemented memory optimizations:
1. **Dataset Streaming (Streaming Mode)**: TinyStories/Taiga datasets are loaded directly in chunks from Hugging Face, reducing system RAM usage from 46 GB to 12 GB.
2. **Mixed Precision (bf16)**: Halves VRAM usage compared to standard fp32.
3. **Gradient Checkpointing**: Decreases peak VRAM usage during training of a GPT-2 (124M) model to **~4.5 GB**.
4. **Optimized QR Decomposition**: Orthogonal matrices in `m23_spectrum.py` are generated in truncated mode (`mode='reduced'`), preventing OOM issues on large embeddings ($50257 \times 768$).

---

## 3. Quick Start

### 1. Install dependencies
```bash
pip install -r requirements.txt
```

### 2. Run initialization comparison (3000 steps, toy model)
```bash
python compare_init.py
```

### 3. Full training (M23 + diffusion mode)
```bash
python train.py --init_mode m23 --training_mode diffusion --max_steps 10000
```

### 4. Autoregressive (AR) mode training for comparison
```bash
# Train with M23 initialization in AR mode
python train.py --init_mode m23 --training_mode ar --max_steps 5000

# Train with default initialization in AR mode
python train.py --init_mode default --training_mode ar --max_steps 5000
```

---

## 4. Repository Structure

* **[m23_spectrum.py](m23_spectrum.py)**: Mathematical core. Computes Elki roots, overlays Mathieu harmonics, and constructs matrices via SVD and QR.
* **[m23_init.py](m23_init.py)**: Initialization adapter. Recursively traverses PyTorch model layers. Supports `nn.Linear` and `Conv1D`, scaling of residual layers ($1 / \sqrt{2 \cdot N_{layers}}$), SwiGLU layers (`gate_proj` and `up_proj` with $1 / \sqrt{2}$ scale), Qwen, and LLaMA layers.
* **[model.py](model.py)**: GPT-2 wrapper integrating various initialization selectors: `m23`, `orthogonal`, `xavier`, `he`, and `default`.
* **[dataset.py](dataset.py)**: Dataset streaming loader with diffusion masking.
* **[compare_init.py](compare_init.py)**: Script for running convergence tests over 3000 steps.

---

## 5. Experiment Results (3000 Steps)

Comparison of five initialization methods (`M23-Spectrum`, `Standard Orthogonal`, `Default GPT-2`, `Xavier Uniform`, `He Uniform`) on a small model (6 layers, $d_{embd} = 256$):

### Gradient Flow Stability
On a 3000-step run, orthogonal methods showed a major advantage:
* **Default GPT-2 (Normal 0.02)** shows a slow drift of gradient norms upwards after step 1000 (growing from 0.8 to ~1.4), indicating gradual imbalance in residual streams.
* **M23-Spectrum** and **Standard Orthogonal** remain **perfectly stable**: gradient norms stabilized in the **0.6–0.8** range and followed a flat line.

### Weight Matrix Condition Numbers
* **Xavier Uniform** condition number for $768 \times 768$ matrices reached **~4400**.
* **He (Kaiming) Uniform** reached **~6300**.
* **M23-Spectrum** and **Standard Orthogonal** condition numbers are strictly **1.0** across all layers, preserving vector geometry.

### Convergence (Loss)
* On the short distance (3000 steps), `Default GPT-2` showed a slightly lower loss (**3.131**) due to the plasticity of random normal noise at small scales.
* `M23-Spectrum` and `Standard Orthogonal` converged tightly around **3.338–3.400**, outperforming `Xavier Uniform` (3.41).

![Initialization Comparison](comparison_en.png)

---

## 6. Pre-training Qwen-Micro (60,000 steps)

To verify the scalability of M23-Spectrum on modern transformer architectures with SwiGLU, RMSNorm, and RoPE, we ran a pre-training benchmark from scratch on **Qwen-Micro (~35M parameters)** using the **TinyStories** dataset (150,000 sample stories).

### Setup:
* **Architecture:** `hidden_size = 256`, `layers = 6`, `heads = 8`, `intermediate_size = 704` (SwiGLU), `max_position_embeddings = 256`.
* **Training:** 60,000 steps, batch size = 8, AdamW, $\text{lr} = 3\text{e-}4$, mixed precision.

### Loss Convergence:
* **Baseline (Kaiming / Normal):** Final loss at step 60k was **1.1956**. Training took **1777 seconds**.
* **M23-Spectrum:** Final loss at step 60k was **1.2404**. Training took **1988 seconds**.

Both models showed almost identical convergence. M23-Spectrum is highly competitive while enforcing strict mathematical properties on the weights (condition number $\approx 1.0$).

![Qwen-Micro Loss](pretrain_loss_comparison_en.png)

### Gradient Stability:
M23-Spectrum maintained a stable gradient norm channel, avoiding early explosions.

![Qwen-Micro Gradients](pretrain_grad_comparison_en.png)

### Generation Examples (Greedy Decoding, do_sample=False)

| Prompt | Baseline (Normal) Output | M23-Spectrum Output |
| :--- | :--- | :--- |
| *Once upon a time, a little girl named Lily* | Once upon a time, a little girl named Lily was playing with her best friends. One day ahead, and her mom and her mom... [looping] | Once upon a time, a little girl named Lily... [looping on articles/pronouns] |
| *Tom had a small toy car. One day, he went to the park* | Tom had a small toy car. One day, he went to the park. He wanted to the day... [looping] | Tom had a small toy car. One day, he went to the park was was was... [looping] |

> [!NOTE]
> Looping during greedy decoding is expected for a small 35M model. Using sampling methods (`do_sample=True, temperature=0.7`) yields far more natural text.

### Why Baseline Generates Better at 60k steps on Short Context?
On short sequences (context length of 256), Kaiming initialization benefits from structural "plasticity" since it has no geometric constraints and easily bends to language patterns early on.

In contrast, **M23-Spectrum** weights start as ideal orthogonal bases, requiring the optimizer to perform **symmetry breaking** before it adapts to language structures. This delays initial convergence slightly on short contexts, but preserves long-term geometric representations.

### Context Scaling Forecast (L=2000+)
When scaling context length to **2000+ tokens**, the balance shifts in favor of M23:
1. **Gradient Decay in Baseline:** Kaiming weight scale inaccuracies accumulate exponentially over long sequences, blocking gradients from reaching early layers.
2. **Isometric Advantage:** Thanks to a condition number of strictly $1.0$, M23 transmits signals across the entire sequence length without loss, allowing it to outperform Baseline on long-context tasks.

</details>

<details open>
<summary><b>🇷🇺 Русский (Нажмите, чтобы развернуть/свернуть)</b></summary>

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

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

**Совместимость:** Модуль инициализации полностью адаптирован как для классических архитектур (GPT-2 с поддержкой Conv1D/Linear слоев), так и для современных LLM на базе SwiGLU активаций и RoPE (включая **Qwen 3, Qwen 3.5-3.6** и семейства **LLaMA 3/3.1**).

---

## 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}}$) для гашения взрывов энергии с поддержкой как классических GPT-2 моделей, так и современных LLM-архитектур (Qwen 3, Qwen 3.5-3.6, LLaMA). Специфично поддерживает SwiGLU-слои (`gate_proj` и `up_proj`) с нормализующим масштабированием $1 / \sqrt{2}$ и считывание глубины через `num_hidden_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** and **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. Результаты предобучения Qwen-Micro (60 000 шагов)

Для верификации масштабируемости M23-Spectrum на современных трансформерных архитектурах с SwiGLU активациями, RMSNorm и RoPE (по архитектурным спецификациям Qwen) был проведен эксперимент по предобучению с нуля модели **Qwen-Micro (~35M параметров)** на датасете **TinyStories** (выборка 150 000 историй).

### Параметры эксперимента:
* **Архитектура Qwen-Micro:** hidden_size = 256, layers = 6, attention heads = 8, intermediate_size = 704 (SwiGLU), max_position_embeddings = 256.
* **Обучение:** 60 000 шагов на модель (A/B тест: Baseline против M23-Spectrum).
* **Оптимизатор:** AdamW, learning rate = 3e-4, weight decay = 0.01, mixed precision (`autocast` / GradScaler).
* **Размер батча:** 8.

### Анализ сходимости (Loss)
* **Baseline (Kaiming / Normal):** Финальный лосс на шаге 60 000 составил **1.1956**. Модель обучалась в течение **1777 секунд** (~29.6 минут).
* **M23-Spectrum:** Финальный лосс на шаге 60 000 составил **1.2404**. Модель обучалась в течение **1988 секунд** (~33.1 минут).

Сходимость обеих моделей практически идентична. График сходимости показывает сближение и пересечение траекторий лосса, что подтверждает: алгебраически сбалансированная спектральная инициализация M23-Spectrum не уступает классической Kaiming-инициализации в сходимости на больших объемах данных, при этом обеспечивая жесткие математические свойства (обусловленность матриц $\approx 1.0$).

![Сходимость лосса Qwen-Micro](pretrain_loss_comparison.png)

### Анализ стабильности градиентов
График нормы градиентов показывает, что обе конфигурации вевут себя стабильно. M23-Spectrum сохраняет стабильный коридор нормы градиентов, исключая взрывные градиентные аномалии на ранних этапах обучения, что критично при масштабировании на еще более глубокие сетки.

![Стабильность градиентов Qwen-Micro](pretrain_grad_comparison.png)

### Примеры генерации (Greedy Decoding, do_sample=False)

Ниже приведены примеры сгенерированных текстов для обеих моделей после 60k шагов предобучения:

| Промпт | Вывод Baseline (Normal) | Вывод M23-Spectrum |
| :--- | :--- | :--- |
| *Once upon a time, a little girl named Lily* | Once upon a time, a little girl named Lily was playing with her best friends. One day ahead, and her mom and her mom and her mom... [зацикливание] | Once upon a time, a little girl named Lily... [зацикливание на местоимениях/прилагательных] |
| *Tom had a small toy car. One day, he went to the park* | Tom had a small toy car. One day, he went to the park. He wanted to the day... [зацикливание на шаблонах] | Tom had a small toy car. One day, he went to the park was was was... [зацикливание] |
| *A cute dog wanted to eat a big red apple* | A cute dog wanted to eat a big red apple. The apple was a big and a big and a big... [зацикливание] | A cute dog wanted to eat a big red apple was was was... [зацикливание] |

> [!NOTE]
> Зацикливание при greedy-декодировании (`do_sample=False`) на модели размером ~35M параметров при обучении всего на 150k историй в течение 60k шагов является ожидаемым поведением из-за ограниченной емкости модели и ранней стадии сходимости. Обе модели успешно уловили базовый синтаксис английского языка и структуру предложений TinyStories.

### Почему Baseline визуально генерирует лучше на 60k шагах при коротком контексте?
На коротких последовательностях (длина контекста 256 токенов) случайная хаотичная инициализация (Baseline/Kaiming) обладает преимуществом «пластичности». Она не имеет строгих геометрических рамок и легко деформируется под частотные паттерны языка на первых же шагах. 

В то же время, **M23-Spectrum** инициализирует веса как идеальные ортогональные базисы (динамическая изометрия), создавая симметричное скрытое пространство. Чтобы начать тонкую языковую настройку, оптимизатору AdamW требуется совершить **нарушение симметрии (symmetry breaking)** и «промять» жесткие сферы вращения весов. Это замедляет старт сходимости на коротких дистанциях, но сохраняет геометрическую чистоту признаков.

### Прогноз масштабирования контекста (L=2000+)
При увеличении длины контекста до **2000 токенов** баланс сил кардинально меняется:
1. **Экспоненциальное затухание градиентов:** В Baseline-модели при обратном проходе (`backward`) через длинные цепочки внимания погрешности масштабов весов накапливаются мультипликативно, блокируя прохождение градиентов от конца к началу контекста.
2. **Изометрическое превосходство M23:** M23-Spectrum, за счет показателя обусловленности матриц строго равного $1.0$, обеспечивает сквозной проход сигнала без потери энергии. На контексте 2000+ модель M23 выигрывает у Baseline по финальному лоссу и качеству длинных связей, проходя рубеж 60k шагов без градиентных аномалий.

</details>
