Andrii Bidochko

• Updated: August 20, 2026
• 6 min read

Dion3: Full‑Stack Orthogonal Updates – A Deep Dive

Direct Answer

Dion3 introduces a full‑stack redesign of the Muon optimizer that slashes orthogonalization overhead by up to six times while preserving—or even improving—model loss. By rethinking the algorithmic core, kernel implementation, and communication pattern, Dion3 makes high‑quality orthogonal updates practical for large‑scale, sharded training workloads.

Background: Why This Problem Is Hard

Training modern deep networks at scale relies on sophisticated optimizers that keep weight matrices well‑conditioned. The Muon optimizer pioneered the use of Newton‑Schulz orthogonalization to maintain orthogonal weight structures, a property that stabilizes training and often yields better generalization. However, the orthogonalization step is cubic in the matrix dimension, turning into a massive FLOP burden as models grow.

When weights are sharded across multiple GPUs or TPUs, the problem compounds: each shard must independently run the Newton‑Schulz iteration, then exchange intermediate results. The resulting communication latency erodes the theoretical speed gains Muon promises, especially in distributed training pipelines where bandwidth is already a scarce resource.

Existing work either accepts the overhead—limiting Muon to small‑to‑medium models—or abandons orthogonal updates altogether, reverting to cheaper but less stable optimizers like Adam or SGD. The gap leaves practitioners without a scalable, high‑performance orthogonal optimizer for today’s trillion‑parameter models.

What the Researchers Propose

The authors present Dion3, a comprehensive overhaul of the Muon stack that attacks the bottleneck on three fronts:

• Gram Newton‑Schulz algorithm: A reformulation that works on the Gram matrix (WᵀW) instead of the weight matrix itself, reducing the asymptotic FLOP count from O(n³) to O(n²).
• CuteDSL kernels: Custom domain‑specific language kernels that exploit the symmetry of the Gram matrix, cutting memory traffic and enabling SIMD‑friendly execution.
• Megabatching strategy: Aggregates orthogonalization work across multiple optimizer steps, amortizing communication costs and allowing a single all‑reduce per megabatch.
• Fractional row orthogonalization: Instead of orthogonalizing the entire momentum matrix each step, Dion3 randomly selects a subset of rows, preserving the statistical benefits while slashing compute.

Collectively, these innovations form a “full‑stack” solution that touches algorithm design, low‑level kernel engineering, and distributed system orchestration.

How It Works in Practice

Conceptual Workflow

During each training iteration, Dion3 follows a four‑stage pipeline:

1. Momentum accumulation: The optimizer gathers gradients into a momentum matrix M, identical to Muon’s approach.
2. Row sampling: A configurable fraction (e.g., 30 %) of M’s rows is chosen for orthogonalization. The remaining rows are carried forward unchanged.
3. Gram Newton‑Schulz update: For the sampled rows, Dion3 computes the Gram matrix G = MᵀM, then applies a Newton‑Schulz iteration on G to approximate its inverse square root. This yields an orthogonal correction that is projected back onto the original rows.
4. Megabatch reduction: Orthogonalization results from several consecutive steps are batched together; a single all‑reduce synchronizes the correction across shards, dramatically lowering network chatter.

Interaction Between Components

The Gram Newton‑Schulz step is the computational heart. By operating on G, the algorithm leverages the fact that G is symmetric positive‑definite, allowing the CuteDSL kernels to unroll the matrix multiplication and reuse intermediate results. The sampled‑row approach ensures that the size of G stays proportional to the fraction of rows, further trimming FLOPs.

After the orthogonal correction, Dion3 updates the original weight matrix W with the corrected momentum, preserving the orthogonal structure without incurring the full cubic cost.

What Makes This Approach Different

• Algorithmic reduction: Moving from W‑centric to Gram‑centric calculations cuts the theoretical complexity.
• Kernel‑level symmetry exploitation: CuteDSL generates code that reuses symmetric entries, halving memory reads.
• Communication amortization: Megabatching turns many small all‑reduces into one larger, more efficient operation.
• Stochastic orthogonalization: Random row selection retains the regularizing effect of orthogonal updates while dramatically lowering per‑step work.

Evaluation & Results

The authors benchmarked Dion3 on three representative workloads:

• Vision Transformer (ViT‑B/16) trained on ImageNet‑1k.
• GPT‑style language model (2.7 B parameters) on a synthetic next‑token task.
• Graph Neural Network (GAT) on the OGB‑Products dataset.

Across all settings, Dion3 achieved:

• Optimizer step time reduction: 3.8× to 6.0× faster than Muon, with the highest speedup observed on the largest sharded GPT model.
• Loss parity or improvement: Final validation loss matched Muon in Vision and GNN tasks, and improved by ~0.4 % on the language model, indicating that the stochastic row selection does not degrade convergence.
• Scalability gains: Communication overhead dropped by up to 70 % thanks to megabatching, enabling near‑linear scaling up to 64 GPUs.

These results demonstrate that Dion3 delivers the theoretical benefits of orthogonal updates without the prohibitive runtime cost that previously limited Muon’s adoption.

Why This Matters for AI Systems and Agents

For engineers building large‑scale AI agents, training speed and stability are non‑negotiable. Dion3’s reduced optimizer step time translates directly into lower cloud compute bills and faster iteration cycles, allowing teams to experiment with deeper architectures or larger batch sizes.

Orthogonal weight updates are especially valuable for recurrent or transformer‑based agents that suffer from exploding/vanishing gradients. By preserving a well‑conditioned weight space, Dion3 can improve the reliability of long‑horizon policy learning in reinforcement‑learning agents or multi‑modal language‑vision systems.

From an infrastructure perspective, the megabatching communication pattern aligns with existing collective‑communication libraries (e.g., NCCL, Gloo), meaning that existing pipelines on the UBOS platform overview can adopt Dion3 with minimal refactoring.

Moreover, the optimizer’s drop‑in compatibility with Muon means that existing codebases—such as those powering AI marketing agents—can upgrade to Dion3 without rewriting training loops, instantly gaining speed and stability benefits.

What Comes Next

While Dion3 marks a substantial leap, several avenues remain open:

• Adaptive row sampling: Dynamically adjusting the fraction of rows based on training dynamics could further balance speed and convergence.
• Mixed‑precision extensions: Integrating bfloat16 or FP8 support in CuteDSL kernels may unlock additional hardware efficiency on next‑gen accelerators.
• Broader model families: Testing Dion3 on diffusion models, large‑scale recommendation systems, and multimodal encoders will validate its generality.
• Tooling integration: Embedding Dion3 into higher‑level orchestration frameworks—such as the Workflow automation studio—could automate megabatch scheduling and row‑sampling policies.

Developers interested in early access can explore the open‑source Dion3 GitHub repository and experiment with the UBOS for startups sandbox environment.

References

Dion3: Full‑Stack Orthogonal Updates (arXiv)

Andrii Bidochko

CTO UBOS

Andrii Bidochko is an AI entrepreneur and researcher focused on AI agents, reinforcement learning, and autonomous systems. He writes about the technologies shaping the future of machine intelligence, from frontier models and agent architectures to real-world AI applications.