MA-RMSNorm: Breaking the Scale-For-Intelligence Trade-off

Magnitude-Aware RMS Normalization for LLMs

This repository contains the complete validation suite and paper for MA-RMSNorm, a normalization method designed to reduce LLM dependence on parameter scaling by preserving magnitude information that standard normalization discards.

Files

File	Description
paper.pdf	Full paper (8 pages, PDF)
paper.html	Full paper (HTML, printable to PDF)
run_all_experiments.py	Complete experiment suite (6 experiments, ~50s to run)
ma_rmsnorm.py	Core implementation + Transformer training framework
output/fig*.png	All 6 experimental figures
output/all_results.json	Numerical results for all experiments
README.md	This file

Quick Start

# Run all 6 experiments (requires: numpy, scikit-learn, matplotlib)
python3 run_all_experiments.py

# Output will be in output/
# - fig1_numerical_utilization.png : Activation distributions
# - fig2_magnitude_mi.png         : Magnitude information recovery
# - fig3_entropy.png              : Attention entropy vs scale
# - fig4_quantization.png         : Quantization robustness
# - fig5_training.png             : Training dynamics
# - fig6_ablation.png             : Ablation study
# - all_results.json              : Machine-readable results

Method Summary

Standard RMSNorm:  y = (x / RMS(x)) * gamma
MA-RMSNorm:        y = (x / RMS(x)) * gamma * s + alpha * W * log(RMS(x))

Where:

• s = 2.0 : Scale factor (expand effective range)
• alpha = 0.1 : Magnitude injection strength
• W : Learned projection (D parameters per norm site)

Key Findings

1. Range expansion: p99 activations increase from 4.1 (RMSNorm) to 20.6 (MA, s=5)
2. No attention collapse: Entropy stays above 99.5% of max even at 8x scale
3. FP16 robustness: MA-RMSNorm outperforms RMSNorm under FP16 quantization
4. Training stability: Identical convergence to baseline across all scale factors
5. Synergistic components: Scale + magnitude path strictly better than either alone

Requirements

• Python 3.10+
• numpy >= 1.20
• scikit-learn >= 1.0
• matplotlib >= 3.5

Citation

If you use this method or build upon this analysis, please cite:

@techreport{ma_rmsnorm2025,
  title  = {MA-RMSNorm: Breaking the Scale-For-Intelligence Trade-off
            via Magnitude-Aware Normalization},
  author = {Anonymous},
  year   = {2025},
  note   = {Independent Technical Report}
}

License

MIT License — feel free to use, modify, and build upon this work.