Implement the Muon step (momentum + Newton–Schulz orthogonalization) for 2D params, falli

Hard

Optimization and Training Dynamics4 / 5

Chapter 074 / 5

Optimization and Training Dynamics

Implement the Muon step (momentum + Newton–Schulz orthogonalization) for 2D params, falli

Hard

Implement the Muon step (momentum + Newton–Schulz orthogonalization) for 2D params, falling back to AdamW for 1D params. [Nebius]

Target interface

newton_schulz(G, steps=5, eps=1e-07)

Complete the starter code in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.

solution.pyPython

Local draft

import torch

def newton_schulz(G, steps=5, eps=1e-07):
    raise NotImplementedError

class Muon:

    def __init__(self, params, lr=0.02, momentum=0.95, wd=0.0, adamw_lr=0.0003, betas=(0.9, 0.999), eps=1e-08):
        raise NotImplementedError

    @torch.no_grad()
    def step(self):
        raise NotImplementedError

Local draft

⌘/Ctrl↵to submit

AI review

Ready when you are

Submit your solution for a score and feedback on correctness, edge cases, and clarity. Any correct approach is accepted.

Chapter 074 / 5

Optimization and Training Dynamics

Implement the Muon step (momentum + Newton–Schulz orthogonalization) for 2D params, falli

Hard

Implement the Muon step (momentum + Newton–Schulz orthogonalization) for 2D params, falling back to AdamW for 1D params. [Nebius]

Target interface

newton_schulz(G, steps=5, eps=1e-07)

Complete the starter code in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.