M23-Spectrum: инициализация весов нейросети через теорию групп Матьё

TL;DR: Вместо случайных весов Xavier/He — детерминистические, математически гарантированные начальные веса на основе алгебры группы . Сходимость в 2.8× быстрее, condition number весовых матриц улучшается в 8× раз, и полный SR-пайплайн до 29–30 dB PSNR на ~900K параметров.

Зачем вообще ду��ать об инициализации весов?

Когда мы говорим об обучении глубоких нейронных сетей, первое, о чём думают — это архитектура, функция потерь, learning rate. Инициализация весов кажется скучной технической деталью: «ну Xavier/He поставил и забыл». Но за этой простотой скрывается фундаментальная проблема.

Xavier (Glorot, 2010) и He (2015) инициализируют веса из случайных распределений с дисперсией, масштабированной под размер слоя. Это работает хорошо для неглубоких сетей, но с ростом глубины возникает системная проблема: спектральный радиус матрицы весов отклоняется от 1, и сигнал либо затухает, либо взрывается при прохождении через десятки слоёв.

Динамическая изометрия — концепция, которая говорит: чтобы сигнал сохранялся, нужно $rho(W) approx 1$ на каждом слое. Добиться этого статистически сложно, особенно стабильно. Но что если взять структуру, где это гарантировано алгебраически?

Группа Матьё M₂₃ и полиномы Элкиса

Группа Матьё $M_{23}$ — одна из 26 спорадических простых групп в классификации конечных простых групп. Она имеет порядок 10 200 960 и обладает исключительно богатой спектральной структурой.

Нас интересует не абстрактная алгебра, а конкретный числовой объект — полином Элкиса для $M_{23}$:

g4+g3+9g2−10g+8=0g4+g3+9g2−10g+8=0

Его четыре корня $g_1, g_2, g_3, g_4 in mathbb{C}$ становятся «семенами» для порождения расширенного спектра через три семейства полиномов:

• $P_2$: $z^2 - gz + 1 = 0$ — квадратичные расширения

• $P_3$: $z^3 + gz - 1 = 0$ — кубические расширения

• $P_4$: $z^4 + gz^3 - (g^2+1)z + g^{-1} - g = 0$ — квартичные расширения

Для каждого из 4 корней мы строим корни всех трёх семейств — итого 40 комплексных компонент. Это и есть базовый M23-спектр.

def _build_m23_base_spectrum() -> np.ndarray:
    elkies = np.roots([1, 1, 9, -10, 8])   # 4 корня полинома Элкиса
    parts  = [elkies]
    for g in elkies:
        parts.append(np.roots([1, -g, 1]))               # P2: 2×4 = 8
        parts.append(np.roots([1, g, 0, -1]))             # P3: 3×4 = 12
        parts.append(np.roots([1, g, 0, -(g**2+1), -g])) # P4: 4×4 = 16
    return np.concatenate(parts)  # 4 + 8 + 12 + 16 = 40

Как из спектра строится матрица весов

Для слоя с fan_in входов нужно расширить 40 компонент до нужного размера. Если fan_in > 40 — тайлинг со сдвигом фазы:

speck=base⋅ei2πk/Ntiles,k=0…Ntiles−1speck=base⋅ei2πk/Ntiles,k=0…Ntiles−1

Это гарантирует равномерное покрытие комплексной плоскости без «провалов» в спектре.

Но��мализация:

λstable=λrawmax⁡∣λraw∣⋅fan_inλstable=max∣λraw∣⋅fan_inλraw

Теперь строим матрицу весов $(fan_out times fan_in)$ с фазовой структурой:

Wrc=λ(r+c)  N⋅ei2πc2/fan_inWrc=λ(r+c)modN⋅ei2πc2/fan_in

Квадратичная фаза $e^{i 2pi c^2 / text{fan_in}}$ — это дискретный чирп-сигнал, обеспечивающий максимальное рассеяние энергии по всем строкам.

Финальный шаг — ортогонализация через SVD:

Wfinal=Re(UV⊤),[U,Σ,V⊤]=SVD(Wcomplex)Wfinal=Re(UV⊤),[U,Σ,V⊤]=SVD(Wcomplex)

В результате все сингулярные значения равны 1 — это и есть динамическая изометрия.

def m23_init_tensor(tensor, variant="orthogonal", seed=None):
    # ... строим матрицу с фазовой структурой ...
    U, _, Vt = np.linalg.svd(mat, full_matrices=False)
    mat = U @ Vt                          # ортогональная матрица
    weights = np.real(mat).reshape(shape).astype(np.float32)
    tensor.copy_(torch.from_numpy(weights))

M23-RLFN: суперрезолюционная сеть

Теория это хорошо, но нужна практическая проверка. Я реализовал M23-RLFN — сеть для задачи super-resolution (апскейл ×4), объединяющую M23-инициализацию с архитектурой победителя NTIRE 2022 Efficient SR Challenge.

Архитектура

LR input (128×128)
  │
  ├─── bicubic(×4) ──────────────────────────────────────┐
  │                                                       │
  └─► Head Conv → [RLFB × 8] → body_conv → Upsample → (+) → SR (512×512)
                      ↑________________________| (body residual)

RLFB (Residual Local Feature Block) — три свёртки 3×3 + ESA + residual connection. Это сердце архитектуры RLFN, доказавшей эффективность при ограниченных параметрах.

ESA (Enhanced Spatial Attention) — механизм внимания с multi-scale max pooling и bilinear upsampling для пространственной фокусировки на деталях.

Global skip connection — итоговое SR-изображение это сумма: SR = detail_net(LR) + bicubic(LR). Сеть учится добавлять детали поверх грубого апскейла, а не генерировать всё изображение с нуля. Это стабилизирует обучение и ускоряет сходимость.

PixelShuffle вместо deconvolution — субпиксельное upsampling без шахматных артефактов.

from m23_sr_engine import M23_RLFN
import torch

model = M23_RLFN(n_feats=52, n_blocks=8, scale=4).cuda()
# Params: ~0.90M

lr = torch.randn(1, 3, 128, 128).cuda()
sr = model(lr).clamp(0, 1)   # → [1, 3, 512, 512]

Функции потерь: почему не MSE

Стандартный L2 (MSE) для SR создаёт «усреднённые» изображения: сеть уходит в размытие, потому что это минимизирует среднеквадратичную ошибку при неопределённости. Я ��спользую комбинацию двух лоссов:

Charbonnier Loss

Lcharb=1N∑i(y^i−yi)2+ε2Lcharb=N1∑i(y^i−yi)2+ε2

Это сглаженный L1: при $varepsilon = 10^{-3}$ ведёт себя как L1 для больших ошибок и как L2 вблизи нуля. Устойчивее к выбросам, чем MSE, и меньше размывает детали.

Frequency Loss (FFT-based)

Идея из CVPR 2024 FDL: штрафуем за несоответствие амплитудных спектров:

Lfreq=λ(L1(∣FFT(y^)∣,∣FFT(y)∣)+0.1⋅1−cos⁡(∠y^−∠y)‾)Lfreq=λ(L1(∣FFT(y^)∣,∣FFT(y)∣)+0.1⋅1−cos(∠y^−∠y))

Высокочастотные компоненты — это рёбра, текстуры, мелкие детали. Пространственный L1 их «не замечает» (вклад мал), а спектральный лосс явно штрафует за их потерю. Фазовая часть гарантирует пространственную согласованность структур.

CombinedSRLoss

L_total = L_charbonnier + 0.05 * L_frequency

Коэффициент freq_weight=0.05 подобран эмпирически: при >0.1 сеть начинает генерировать артефакты рингинга.

Тренировочный пайплайн на DIV2K

DIV2K — стандартный бенчмарк для SR: 800 тренировочных + 100 валидационных HR-изображений высокого качества. LR-пары генерируются на лету через bicubic downsampling с антиалиасингом:

lr = F.interpolate(hr.unsqueeze(0), size=(lr_size, lr_size),
                   mode="bicubic", align_corners=False, antialias=True)

Аугментации

• Random crop 256×256 из HR (64×64 для LR при ×4)

• Random horizontal/vertical flip (p=0.5 каждый)

• Random rotation 90°/180°/270°

Multi-stage Warm-Start

Стратегия из оригинальной статьи RLFN:

На каждой новой стадии оптимайзер перезапускается с новым LR — это позволяет «вырваться» из локальных минимумов.

AMP + Gradient Clipping

Обучение с torch.amp.autocast на RTX 4070 Ti Super даёт примерно 1.6× ускорение по сравнению с FP32. Клиппинг градиентов (max_norm=1.0) предотвращает взрыв градиентов на ранних итерациях.

python train_div2k.py 
    --scale 4 --n_feats 52 --n_blocks 8 
    --total_iters 300000 --batch_size 16 
    --freq_weight 0.05 --val_every 5000

Результаты и бенчмарки (предварительные)

Супер-разрешение Set5 ×4

M23-RLFN конкурирует с RLFN при вдвое большем количестве параметров — есть куда расти с точки зрения архитектурной эффективности, но M23-инициализация уже сейчас даёт хорошую стартовую точку.

Сходимость (24-слойная сеть, DIV2K)

Condition number (отношение max/min сингулярного значения) — ключевая метрика здоровья весовых матриц. При 145 у He-инициализации матрица близка к сингулярной, что ухудшает обусловленность задачи оптимизации.

Производительность (RTX 4070 Ti Super, 16GB)

Инференс одного изображения

pythonfrom train_div2k import upscale_image

upscale_image(
    model_path = "checkpoints/m23_rlfn/m23_rlfn_best.pth",
    input_path  = "photo_lr.png",
    output_path = "photo_4x.png",
    scale=4,
    device="cuda",
)

from train_div2k import upscale_image upscale_image( model_path = "checkpoints/m23_rlfn/m23_rlfn_best.pth", input_path = "photo_lr.png", output_path = "photo_4x.png", scale=4, device="cuda", )

Ограничения и планы

Текущие ограничения:

• Инициализация слоёв >10K нейронов: SVD занимает 0.3–1s. Это разовая операция при старте обучения, но неудобно при частом переинициализации.

• Для архитектур с weight sharing (например, tied embeddings в трансформерах) рекомендуется variant="scaled" вместо "orthogonal".

• ~900K параметров пока уступают эффективности RLFN-S при 317K по соотношению PSNR/параметры.

Roadmap (v3.1+):

• GPU-ускорение инициализации через CUDA batched SVD (torch.linalg.svd)

• Перцептивный loss (LPIPS) для фотореалистичного SR

• Поддержка ×2 и ×3 апскейла

• TensorRT и ONNX экспорт для real-time инференса

• Интеграция в игровые движки (шейдер-пайплайн)

Установка и репозиторий

bashgit clone https://github.com/m23spectrum/m23-spectrum.git
cd m23-spectrum
pip install torch torchvision pillow numpy

git clone https://github.com/m23spectrum/m23-spectrum.git cd m23-spectrum pip install torch torchvision pillow numpy

Репозиторий: github.com/m23spectrum/m23-spectrum

Лицензия: MIT.

Выводы

M23-Spectrum — это попытка ответить на вопрос: «а что если инициализацию весов строить не статистически, а алгебраически?». Полиномы Элкиса для группы $M_{23}$ задают спектральную структуру с гарантированными свойствами изометрии, которые сохраняются через SVD-ортогонализацию.

Практический результат: на задаче SR с архитектурой M23-RLFN метод даёт 2.8× ускорение сходимости и 8.2× улучшение condition number весовых матриц по сравнению с He-инициализацией.

Проект находится в стадии активной разработки. Буду рад issues, pull requests и критике методологии в комментариях.