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

pip install -r requirements.txt

2. Run initialization comparison (3000 steps, toy model)

python compare_init.py

3. Full training (M23 + diffusion mode)

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

4. Autoregressive (AR) mode training for comparison

# 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: Mathematical core. Computes Elki roots, overlays Mathieu harmonics, and constructs matrices via SVD and QR.
• 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: GPT-2 wrapper integrating various initialization selectors: m23, orthogonal, xavier, he, and default.
• dataset.py: Dataset streaming loader with diffusion masking.
• 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).

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$).

Gradient Stability:

M23-Spectrum maintained a stable gradient norm channel, avoiding early explosions.

Generation Examples (Greedy Decoding, do_sample=False)

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.

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. Установка зависимостей

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

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$).

Анализ стабильности градиентов

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

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

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

Почему 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 шагов без градиентных аномалий.