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

Hard

Optimization & Training Dynamics4 / 5

Chapter 074 / 5

Optimization & 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)

Implement the function/class skeleton 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

⌘/Ctrl↵to submit

AI review

No review yet

Submit your solution for a verdict, a score out of 100 and concrete feedback. Any correct approach passes.

Chapter 074 / 5

Optimization & 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)

Implement the function/class skeleton in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.