Add muon optimizer

[Adhithya-Laxman]

Description

This PR implements the Muon optimizer using pure NumPy, completing the neural network optimizers module for the repository.

Muon is a cutting-edge optimizer specifically designed for hidden layer weight matrices in neural networks, using Newton-Schulz matrix orthogonalization iterations for improved convergence and computational efficiency.

This PR addresses part of issue #13662 - Add neural network optimizers module to enhance training capabilities

What does this PR do?

• Implements Muon optimizer that applies Newton-Schulz orthogonalization to gradient matrices
• Specifically optimized for 2D weight matrices (hidden layers)
• Provides superior convergence properties compared to standard optimizers
• Uses momentum with orthogonalized gradient updates
• Memory-efficient implementation with minimal state tracking
• Educational implementation suitable for learning advanced optimization techniques

Implementation Details

• Algorithm: Muon (Matrix-aware Optimizer with Orthogonalization)
• Core Mechanism: Newton-Schulz iterations for gradient orthogonalization
• Update rule:

  ortho_grad = NewtonSchulz(gradient, steps=5)
  velocity = momentum * velocity + ortho_grad
  param = param - learning_rate * velocity
  

• Newton-Schulz Iteration:

  For k steps:
    A = 1.5 * X - 0.5 * X @ (X.T @ X)
  Returns orthogonalized matrix
  

• Key Features:
  • Matrix-aware optimization
  • Spectral norm-based updates
  • Implicit regularization through orthogonalization
  • Momentum accumulation with orthogonalized gradients
• Pure NumPy: No PyTorch, TensorFlow, or other frameworks required
• Educational focus: Clear implementation showing matrix orthogonalization concepts

Why Muon?

Muon represents state-of-the-art optimizer research and offers several advantages:

• Computational Efficiency: ~2x more efficient than AdamW in large-scale training
• Better Convergence: Particularly effective for training deep networks
• Memory Efficient: Maintains only momentum buffer (no second-order moments)
• Geometric Awareness: Exploits matrix structure of neural network parameters
• Recent Innovation: Used in current training speed records for NanoGPT and CIFAR-10

Features

✅ Complete docstrings with parameter descriptions
✅ Type hints for all function parameters and return values
✅ Doctests for correctness validation
✅ Usage example demonstrating optimizer on matrix optimization
✅ PEP8 compliant code formatting
✅ Newton-Schulz orthogonalization implementation
✅ Configurable hyperparameters (learning rate, momentum, iteration steps)
✅ Pure NumPy - no external deep learning frameworks

Testing

All doctests pass:

python -m doctest neural_network/optimizers/muon.py -v

Linting passes:

ruff check neural_network/optimizers/muon.py

Example output demonstrates proper optimization behavior on matrix parameters.

References

• Muon: An optimizer for hidden layers (Original Blog Post)
• Muon GitHub Repository
• Muon is Scalable for LLM Training (arXiv)
• Practical Efficiency of Muon for Pretraining (arXiv)
• Newton-Schulz Iteration - Matrix Functions

Relation to Issue #13662

This PR completes the optimizer sequence outlined in #13662:

• ✅ Stochastic Gradient Descent (SGD) - feat: add SGD optimizer for neural networks #13671
• ✅ Momentum SGD - Add Momentum SGD optimizer implementation #13680
• ✅ Nesterov Accelerated Gradient (NAG) - Added Nesterov and Adam Optimizers #13718
• ✅ Adagrad - Add Adagrad optimizer implementation in Pure Numpy #13681
• ✅ Adam - Added Nesterov and Adam Optimizers #13718
• ✅ Muon (this PR) - FINAL OPTIMIZER

With this PR, the neural network optimizers module is now complete with 6 fundamental optimizers covering classical to cutting-edge optimization techniques.

Use Cases

Muon is particularly effective for:

• Training large language models and transformers
• Deep convolutional networks
• Scenarios requiring computational efficiency at scale
• Research into matrix-aware optimization
• Cases where AdamW shows diminishing returns

Checklist

• I have read CONTRIBUTING.md
• This pull request is all my own work -- I have not plagiarized
• I know that pull requests will not be merged if they fail the automated tests
• This PR only changes one algorithm file
• All new Python files are placed inside an existing directory
• All filenames are in all lowercase characters with no spaces or dashes
• All functions and variable names follow Python naming conventions
• All function parameters and return values are annotated with Python type hints
• All functions have doctests that pass the automated testing
• All new algorithms include at least one URL that points to Wikipedia or another similar explanation

Summary

This PR marks the completion of the neural network optimizers module, providing educators and learners with a comprehensive collection of optimization algorithms from fundamental SGD to cutting-edge Muon. The module now serves as a complete educational resource for understanding neural network training optimization.

---

This PR along with the following PRs collectively fixes issue #13662:

Related PRs:

• feat: add SGD optimizer for neural networks #13671
• Add Momentum SGD optimizer implementation #13680
• Add Adagrad optimizer implementation in Pure Numpy #13681
• Added Nesterov and Adam Optimizers #13718
• This PR (Muon optimizer) - Final optimizer completing the module

Fixes #13662 (This PR completes the neural network optimizers module)

[algorithms-keeper]

Multiple Pull Request Detected

@Adhithya-Laxman, we are extremely excited that you want to submit multiple algorithms in this repository but we have a limit on how many pull request a user can keep open at a time. This is to make sure all maintainers and users focus on a limited number of pull requests at a time to maintain the quality of the code.

This pull request is being closed as the user already has an open pull request. Please focus on your previous pull request before opening another one. Thank you for your cooperation.

User opened pull requests (including this one): #13720, #13718, #13681, #13680, #13646