PaperlayerView exact revision on arXiv

Source-linked AI summary

RODE: A Radial-Orthogonal Decoupled Engine for Optimization

Guoxiang Xu, Bince Qu, Qi Sun, Cheng Zhuo

arXiv:2608.21024v1cs.LG

TL;DR

Additive matrix updates couple norm and direction, making later angular motion depend on uncontrolled scale evolution. RODE separates these coordinates with independent rules and step sizes, combining scalar radius updates with tangent-space Newton–Schulz conditioning. Across the reported language and vision comparisons, RODE outperforms both Muon variants and shows more controlled norm growth, including at 1.5B and 9B scale.

  • Problem

    Additive matrix updates jointly change norm and direction, although norm affects later angular motion and directional learning can drive norm growth.

  • Method

    RODE treats matrix Frobenius norm and conditioned direction as separate coordinates with independent update rules and step sizes.

  • Results

    RODE outperforms both Muon variants across the main language-modeling and vision comparisons and attains the highest mean on GSM8K and MATH-500 among six optimizers.

  • Takeaways & Limitations

    Explicit norm control is supported as an effective mechanism for matrix optimization within the evaluated tasks and settings.

  • Takeaways & Limitations

    RODE applies its radial–directional update only to matrix-valued parameters, adds computational overhead, and has been evaluated on a limited set of architectures and training budgets.

Abstract

from arXiv · show

Modern neural network training increasingly uses matrix-aware optimizers, yet their conditioned matrix step is typically added directly to the weight, jointly changing its norm and direction. This interaction matters because the current norm determines angular motion, while directional learning can drive norm growth and thereby alter later steps. We introduce RODE, which gives the radial and directional components separate update rules and step sizes. RODE explicitly updates the matrix Frobenius norm through a scalar radial rule, while its directional channel performs Newton--Schulz-conditioned updates in the tangent space. Controlled GPT-2 interventions show gains from both direct norm control and RODE's directional update. Across two language-modeling and two image-classification tasks, RODE outperforms both Muon variants in every direct comparison and ends with lower full-model norms. At 1.5B scale, using the learning rate transferred directly from the Qwen2-style LM sweep, RODE lowers loss from 4.145 to 3.346 and final global norm from 11964 to 2183 relative to Muon RMS, with fixed-radius RODE improving further. For Qwen3.5-9B full-parameter fine-tuning, all six optimizers use the same tuning budget and the same formal-training and evaluation settings; RODE outperforms both Muon variants on all four evaluation tasks and attains the highest mean on GSM8K and MATH-500. Thus, decoupling radial and directional dynamics offers a more effective and controllable approach to matrix optimization.

1 Introduction

RODE addresses the coupling between matrix norm and direction in additive updates by assigning radial and directional dynamics separate rules and learning rates. It retains matrix-aware conditioning while making these effects independently controllable.

  • 1 Introduction: Additive matrix updates jointly alter Frobenius norm and direction, although these components have distinct optimization roles.In scale-invariant blocks, rescaling can preserve the represented function while changing later angular motion; elsewhere, scale can affect the function directly.
  • 1 Introduction: Directional learning can increase norm, reducing the angle realized by later absolute steps; Muon generally mixes both effects in one conditioned update.Weight decay opposes growth only through an optimizer- and schedule-dependent balance.
  • 1 Introduction: RODE gives radial and directional signals different update rules and learning rates while retaining Newton–Schulz conditioning for direction.Its directional channel uses tangent projections and a spherical realization, while a scalar rule updates radius.
  • 1 Introduction: RODE evaluates component effects through controlled interventions and overall performance across four language and vision tasks, with additional studies at 1.5B and 9B scale.The evaluation includes hyperparameter transfer at 1.5B scale and full-parameter fine-tuning at 9B scale.

2 Related Work

RODE connects matrix-conditioned update directions with explicit parameter-scale control. Unlike prior methods that normalize, constrain, or separately parameterize scale, it exposes a matrix-wide radius alongside tangent-projected conditioned direction updates.

  • 2 Related Work: RODE connects matrix-valued update directions with parameter-scale control, two lines of work often treated separately.Its design combines matrix conditioning with an explicit radial coordinate.
  • 2 Related Work: Prior matrix-aware optimizers include coordinate-, factor-, eigenbasis-, and momentum-based conditioning, while Muon applies short Newton–Schulz iterations to momentum.These methods primarily concern the update direction or its conditioning.
  • 2 Related Work: Scale-oriented methods expose magnitude, regulate norm growth, or scale additive steps, including weight normalization, AdamP, LARS/LAMB, Fromage, and Nero.For normalized weights, norm growth can slow subsequent directional change, while weight decay can counteract it.
  • 2 Related Work: Hyperball and Magnitude–Direction Decoupling use fixed-norm or separately parameterized scale components, whereas RODE learns one matrix-wide radius and tangent-projects conditioned momentum.RODE therefore differs from fixed-radius and row/column-gain formulations.

3 Method

RODE decomposes each managed matrix into radius and unit-Frobenius direction, then updates the two coordinates separately. Its directional engine reprojects momentum and applies tangent-preserving Newton–Schulz conditioning before a spherical step.

  • 3 Method: RODE represents a managed matrix as a Frobenius radius multiplied by a unit-Frobenius direction, with other parameters using AdamW.The direction is the matrix-conditioned coordinate, while the radius is updated separately.
  • 3 Method: RODE omits the current-norm factor from the directional signal so ηdir directly controls directional step magnitude.A positive scalar does not change the normalized polar direction used by the Newton–Schulz conditioner.
  • 3 Method: The radial channel updates only the matrix Frobenius norm using the current gradient component parallel to the direction.Reprojected historical momentum contributes no radial component, so the radial signal is determined only by the current gradient.
  • 3 Method: RODE reprojects historical momentum onto the current tangent space before adding the current gradient, then forms the tangent momentum for conditioning.Because the direction changes, radial content retained in the raw buffer can re-enter a future tangent space; Appendix D bounds this residual.
  • 3 Method: Projected Newton–Schulz iterations condition tangent momentum, with projection after every polynomial step to remove newly introduced radial components.The resulting conditioned tangent kernel supplies the directional update, whose output is renormalized for numerical stability.
  • 3 Method: RODE-SNR optionally damps directional steps using a running estimate of tangent-gradient scale, while default RODE sets the damping factor to one.The optional variant is separately tuned.

4 Theory

RODE’s theory analyzes radial descent, tangent descent, curvature, and raw-buffer leakage under explicit regularity and alignment assumptions. A separate aligned first-order model yields linear convergence to an η-controlled neighborhood.

  • 4 Theory: The main analysis assumes smoothness, bounded stochastic noise and iterates, tangent-gradient regularity, bounded rotations, and positive conditioning alignment.It applies in the interior regime where the radial floor is inactive and the tangent kernel is nonzero.
  • 4 Theory: An exact displacement expansion separates radial descent, tangent descent, geodesic curvature, and cross-step leakage from raw momentum.A Lyapunov function combines objective value with tangent-momentum tracking error.
  • 4 Theory: The convergence bounds cover both gradient components even though stored raw momentum need not remain perfectly tangent as the direction changes.The error terms depend on radial and directional step sizes, stochastic-noise levels, and assumption constants.
  • 4 Theory: The first-order approximation assumes bounded effective updates, positive alignment with the full gradient, and a Polyak–Lojasiewicz objective condition.Under 0 < 2ηαµ < 1, the model converges linearly to a neighborhood controlled by η.

5 Experiments

RODE is evaluated through controlled component interventions, cross-task benchmarks, scale transfer, large-model fine-tuning, and systems measurements. It consistently improves optimization outcomes over Muon variants while producing more controlled full-model norms, although its radial coordinate adds computational cost.

  • Experimental setup: RODE is evaluated on language modeling and image classification, including GPT-2, Qwen2-style models, ResNet-50, and Qwen3.5-9B fine-tuning.Comparisons use three seeds and within-task protocols; the 1.5B study transfers rates from the smaller Qwen2-style model without another sweep.
  • Optimization quality: 0.0491 and 0.0454 lower loss, plus 2.36 and 4.48 accuracy-point gains, are reported against the stronger Muon variant on GPT-2, Qwen2-style LM, CIFAR-100, and ImageNet-1K, respectively.RODE leads on CIFAR-100, while SOAP leads GPT-2 and RODE-SNR leads the Qwen2-style LM and ImageNet-1K.
  • Optimization quality: Muon RMS finishes at 2.5–5.7× the norm of RODE, while RODE variants attain the lowest observed final full-model norms across the tasks.Full-model norms include all trainable parameters, including tensors delegated to AdamW, and are comparable only within a task.
  • Norm–direction interventions: Fixing Muon-managed radii lowers GPT-2 validation loss from 3.027 to 3.008, and replacing its additive direction update lowers it further to 2.977.Full RODE with learned radii reaches the same 2.977 level, isolating measurable contributions from norm control and tangent directional updates.
  • Scale transfer: At 1.5B scale, transferred rates give RODE lower loss and global norm than Muon RMS, while fixed-radius RODE improves both measures further.The experiment runs for 20000 FineWeb steps and freezes managed matrix radii only for the fixed-radius variant; AdamW-delegated parameters can still change the global norm.
  • Full-parameter fine-tuning: RODE outperforms both Muon variants on all four Qwen3.5-9B evaluation tasks and achieves the highest mean on GSM8K and MATH-500.Adafactor leads on MMLU-STEM and MMLU-Pro Math; all six optimizers receive the same tuning budget and formal-training and evaluation settings.
  • Components and systems cost: Removing projected Newton–Schulz conditioning or directional learning substantially degrades Qwen2-style-LM performance, while RODE stores 611 MB and reaches 0.72× AdamW throughput after fusion.On ResNet-50, fused RODE stores 102.4 MB and reaches 0.64× AdamW throughput, versus 0.70× for Muon RMS.

6 Discussion

RODE makes matrix radius and direction separate optimizer choices, aiming to prevent unintended norm growth from silently governing directional progress. Across the paper’s evaluated settings, it improves optimization quality while reducing observed norms, with compute overhead and scope limitations remaining.

  • RODE combines a learned matrix-radius policy, a separately tuned directional policy, and matrix conditioning retained for direction.
  • Explicit radial control prevents norm evolution from being an uncontrolled by-product of directional learning, while keeping the choice visible and tunable.
  • RODE reduces observed norm growth while improving over both Muon variants in every main language-modeling and vision comparison.
  • At 9B scale, RODE outperforms both Muon variants on all four evaluation tasks and leads the six-optimizer comparison on GSM8K and MATH-500.
  • RODE currently applies radial-directional updates only to matrix-valued parameters, while remaining parameters are delegated to AdamW.
  • The study covers a limited set of language and vision architectures and training budgets; broader scales and model families remain to be explored.

7 Conclusion

RODE separates norm evolution from directional progress by treating matrix Frobenius radius and conditioned direction as distinct coordinates. Experiments support this design across language modeling, vision, scale transfer, and 9B fine-tuning.

  • Direct additive matrix updates couple weight-norm evolution with directional progress, allowing norm growth to alter subsequent optimization dynamics.
  • RODE separates the matrix Frobenius norm and matrix-conditioned direction into distinct optimization coordinates with independent update rules and step sizes.
  • RODE outperforms both Muon variants on all four evaluation tasks at Qwen3.5-9B scale and leads the six-optimizer comparison on GSM8K and MATH-500.
  • RODE’s update takes parameter Wt, gradient gt, raw buffer Bt−1, separate radial and directional learning rates, momentum, Newton–Schulz steps, and safeguards as inputs.
  • The directional update projects the conditioned momentum into the tangent space, applies Newton–Schulz conditioning, and realizes motion on the sphere before renormalization.

A.2 Shape-batched implementation

Shape batching removes the per-matrix Python loop while preserving RODE’s per-matrix computation. The resulting implementation improves throughput but remains slower than Muon RMS in the reported benchmark.

  • The fused implementation stacks same-shaped matrices, gradients, and momentum buffers so the engine operates as batched tensor operations.
  • The fused polar split computes each matrix’s Frobenius radius and normalized direction over the final two dimensions.
  • The batched directional path projects momentum into the tangent space and applies Newton–Schulz conditioning through batched matrix multiplications.
  • 0.59× to 0.72× AdamW: batching raises Qwen2-style-LM throughput, compared with 0.77× for Muon RMS at equal state memory.
  • Matrices of different shapes form separate buckets, and loop and batched implementations agree to bf16 numerical tolerance.

B Optimizer Definitions and Hyperparameters

The paper specifies optimizer routing, sweep and transfer protocols, model-specific training settings, diagnostics, and evaluation procedures. Comparisons use controlled budgets and task-specific protocols documented across the benchmark tables and appendices.

  • RODE’s matrix learning rate is the angular rate ηdir, while ηρ controls its radial update and η1d controls non-matrix parameters.
  • Matrix weight decay is the task-level –wd, while non-matrix AdamW uses weight decay 9 × 10−4 unless specified otherwise.
  • All optimizers within each sweep group share the listed trial count, horizon, data source, model configuration, evaluation protocol, and log-uniform search spaces.
  • The added baselines are swept on Qwen2-style-LM/WikiText-103 and then transferred to standard GPT-2 and CIFAR-100.
  • The Qwen2-style LM benchmark runs 4000 steps with 114M parameters, sequence length 512, effective batch size 128, 200 warmup steps, and cosine decay.
  • Qwen3.5-9B fine-tuning gives every optimizer 20 trials of 800 steps from the same checkpoint and uses the minimum-validation-loss checkpoint for downstream evaluation.
  • The 1.5B transfer experiment reuses smaller-model learning rates without a new sweep and runs three seeds for 20000 steps.

C.4 Qwen3.5-9B tuning and training diagnostics

The supplied diagnostics cover Qwen3.5-9B tuning, GPT-2/WikiText-103 dynamics, and ResNet-50 training behavior across CIFAR-100 and ImageNet-1K. They report validation or training losses, accuracy, gradients, update-to-weight ratios, and full-model norms using multi-seed summaries.

  • Qwen3.5-9B tuning and training diagnostics: Qwen3.5-9B diagnostics compare six optimizers through validation loss during 3000-step formal runs and cumulative best objectives across 20 800-step Optuna trials.Both panels use raw, unsmoothed values and three-seed mean ± sample SD for the formal runs.
  • GPT-2/WikiText-103: GPT-2/WikiText-103 diagnostics track training loss, update step size, and gradient norm across three seeds.The GPT-2 weight norm is reported separately in the main text.
  • Full-model norms: Full-model weight norms on CIFAR-100 and ImageNet-1K are shown as three-seed mean ± sample SD, including parameters delegated to AdamW.RODE directly updates managed-matrix radii.
  • ResNet-50/CIFAR-100: ResNet-50/CIFAR-100 diagnostics include training accuracy, training loss, gradient norm, and update-to-weight ratio.Additional curves complement the reported test-loss and full-model-norm plots.
  • ResNet-50/ImageNet-1K: ResNet-50/ImageNet-1K diagnostics include training accuracy, training loss, gradient norm, and update-to-weight ratio, while main validation results appear in Table 1 and Figure 4.The curves are three-seed means with sample-standard-deviation bands.

D Detailed Global Convergence Proof

The appendix analyzes RODE’s radial and tangent dynamics through stochastic-gradient decompositions, projected Newton–Schulz updates, curvature terms, and raw-buffer leakage. Under stated interior-regime assumptions, a Lyapunov argument yields finite-time and asymptotic convergence-neighborhood bounds.

  • Detailed Global Convergence Proof: RODE’s convergence analysis separates radial descent, tangent descent, geodesic curvature, and cross-step leakage from the raw momentum buffer.A Lyapunov function combines the objective with tangent-momentum tracking error.
  • Gradient decomposition: The stochastic gradient and full gradient are decomposed into radial and tangent components before conditional-expectation analysis.The radial channel uses a guarded scalar update, while tangent momentum is projected as the direction evolves.
  • Assumptions and branch conditions: The proof analyzes the nondegenerate branch with inactive radial floors, bounded states and noise, local tangent-gradient regularity, and aligned projected Newton–Schulz kernels.The radial-floor-active branch requires a different projected-gradient inequality and is not claimed in this appendix.
  • Projected directional engine: Tangent reprojection follows each quintic Newton–Schulz step, removing the component parallel to the current direction before the next step.The resulting projected direction is assumed to retain nondegenerate norm and alignment with the tangent momentum.
  • Convergence neighborhood: Theorem 1 states asymptotic upper bounds for time-averaged radial gradients, normalized tangent momentum, and tangent tracking error under sufficiently small radial learning rates.The finite-time bound follows from the Lyapunov lower bound and nonnegative descent terms; taking lim sup removes the remaining finite-horizon term.
  • Scope of the first-order model: The first-order model is separate from the exact raw-buffer recurrence and requires finite curvature, unequal coordinate stepsizes, radial-floor, and finite-kernel effects to be controlled as residuals.Within its stated scope, the model yields linear convergence to a stepsize-controlled neighborhood under boundedness, alignment, smoothness, and PL assumptions.

E.1 Additional assumptions for PL analysis

The PL analysis adds assumptions for bounded effective updates, local reprojected tangent-gradient regularity, polar stability, and strong directional alignment. Under these conditions, alignment supports a linear-convergence recursion for the first-order model, but not automatically for every finite-step implementation branch.

  • Additional assumptions for PL analysis: The PL setting assumes an L-smooth objective satisfying the Polyak–Łojasiewicz condition, bounded effective update directions, and locally regular reprojected tangent gradients.The local condition accounts for both gradient variation and drift between moving tangent spaces.
  • Additional assumptions for PL analysis: The analysis also assumes bounded tangent gradients, local polar-map stability, and strong directional alignment.The polar-stability condition excludes degenerate points where the polar factor is unstable.
  • Scope of alignment claims: The sufficient alignment lemmas do not establish the condition for every finite-step implementation branch; Theorem 2 assumes effective alignment directly.This limits the scope of the PL argument to the aligned first-order model.
  • Alignment synthesis: Radial and tangent alignment bounds combine because the two subspaces are orthogonal, yielding the full effective-direction alignment required by the PL recursion.Lemma 3 synthesizes the componentwise alignment results into the global condition.
  • Tangent alignment: Projected-polar tangent steps remain positively aligned with the true tangent gradient at sufficiently small learning rates, including scalar damping bounded below by a positive constant.The tangent contribution is tangent by construction after projection.
Loading 2608.21024v1…