Принцип совместимости с симметрией для проектирования оптимизатора: встраивания, головы LM, SwiGLU-многослойные перцептроны и маршрутизаторы MoE

Аннотация

Давно сохраняется заметное геометрическое несоответствие в практике глубокого обучения. В то время как современные архитектуры нейронных сетей естественным образом обладают богатыми свойствами симметрии и эквивариантности, популярные оптимизаторы, такие как Adam и его варианты, работают по своей сути покоординатно, что не позволяет им учитывать структуры эквивариантности пространства параметров. Мы устраняем это несоответствие, вводя принцип проектирования оптимизаторов, совместимый с симметрией: правило градиентного обновления должно быть эквивариантным относительно группы симметрии, действующей на соответствующий весовой блок. Следуя этому принципу, мы сначала предлагаем единую перспективу на биортогонально эквивариантные обновления для общих матричных слоёв, используемые в методах стохастического спектрального спуска, Muon, Scion и полярных градиентных методов. Что более важно, переходя от ортогональных групп к симметриям перестановок и общих сдвигов, мы выводим совместимые с симметрией оптимизаторы для блоков параметров, чьи симметрии отличаются от симметрий общих матричных слоёв: матриц вложений и LM-головок, проекций SwiGLU MLP и матриц маршрутизации MoE. Эти конструкции включают односторонние спектральные, строчно-нормированные, гибридные строчно-нормированные/спектральные, строчно-чувствительные, столбцово-чувствительные, центрированные строчно-нормированные и лево-спектральные обновления. Они образуют сквозной послойный стек оптимизаторов, в котором каждому основному классу матричных параметров назначается обновление, чья эквивариантность соответствует его группе симметрии. Мы подтверждаем этот принцип с помощью экспериментов по предварительному обучению на плотных и разреженных языковых моделях MoE, включая архитектуры в стиле Qwen3-0.6B, Gemma 3 1B, OLMoE-1B-7B и уменьшенные gpt-oss. Во всех этих экспериментах обновления, совместимые с симметрией, последовательно улучшают итоговое значение функции потерь на валидации, а в ряде случаев и стабильность обучения, по сравнению с соответствующими обновлениями AdamW.

English

A striking geometric disparity has long persisted in the practice of deep learning. While modern neural network architectures naturally exhibit rich symmetry and equivariance properties, popular optimizers such as Adam and its variants operate inherently coordinate-wise, rendering them unable to respect the equivariance structures of the parameter space. We address this disparity by introducing a symmetry-compatible principle for optimizer design: the gradient update rule should be equivariant under the symmetry group acting on the corresponding weight block. Following this principle, we first provide a unified perspective on bi-orthogonally equivariant updates for general matrix layers, as employed by stochastic spectral descent, Muon, Scion, and polar gradient methods. More importantly, by moving from orthogonal groups to permutation and shared-shift symmetries, we derive symmetry-compatible optimizers for parameter blocks whose symmetries differ from those of general matrix layers: embedding and LM head matrices, SwiGLU MLP projections, and MoE router matrices. These constructions include one-sided spectral, row-norm, hybrid row-norm/spectral, row-aware, column-aware, centered row-norm, and left-spectral updates. They yield an end-to-end layerwise optimizer stack in which each major matrix-valued parameter class is assigned an update whose equivariance matches its symmetry group. We corroborate this principle through pre-training experiments on dense and sparse MoE language models, including Qwen3-0.6B-style, Gemma 3 1B-style, OLMoE-1B-7B-style, and downsized gpt-oss architectures. Across these experiments, symmetry-compatible updates consistently improve final validation loss, and in several cases training stability, over corresponding AdamW updates.