neural_network.optimizers.muon_optimizer

Muon Optimizer

Implements Muon optimizer for neural network hidden layers using NumPy. Muon uses Newton-Schulz orthogonalization iterations for improved convergence.

Reference: https://kellerjordan.github.io/posts/muon/ Author: Adhithya Laxman Ravi Shankar Geetha Date: 2025.10.21

Attributes

optimizer

Classes

Muon	Muon optimizer for hidden layer weight matrices.

Module Contents

class neural_network.optimizers.muon_optimizer.Muon(learning_rate: float = 0.02, momentum: float = 0.95, ns_steps: int = 5)

Muon optimizer for hidden layer weight matrices.

Applies Newton-Schulz orthogonalization to gradients before updates.

newton_schulz_orthogonalize(matrix: numpy.ndarray) → numpy.ndarray

Orthogonalize matrix using Newton-Schulz iterations.

Args:

matrix (np.ndarray): Input matrix.

Returns:

np.ndarray: Orthogonalized matrix.

>>> optimizer = Muon()
>>> mat = np.array([[1.0, 0.5], [0.5, 1.0]])
>>> orth = optimizer.newton_schulz_orthogonalize(mat)
>>> orth.shape
(2, 2)

update(param_id: int, params: numpy.ndarray, gradients: numpy.ndarray) → numpy.ndarray

Update parameters using Muon.

Args:

param_id (int): Unique identifier for parameter group. params (np.ndarray): Current parameters. gradients (np.ndarray): Gradients of parameters.

Returns:

np.ndarray: Updated parameters.

>>> optimizer = Muon(learning_rate=0.1, momentum=0.9)
>>> params = np.array([[1.0, 2.0], [3.0, 4.0]])
>>> grads = np.array([[0.1, 0.2], [0.3, 0.4]])
>>> updated = optimizer.update(0, params, grads)
>>> updated.shape
(2, 2)

learning_rate = 0.02

momentum = 0.95

ns_steps = 5

velocity: dict[int, numpy.ndarray]

neural_network.optimizers.muon_optimizer.optimizer