import%20marimo%0A%0A__generated_with%20%3D%20%220.24.0%22%0Aapp%20%3D%20marimo.App()%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20time%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20pandas%20as%20pd%0A%20%20%20%20import%20plotly.graph_objects%20as%20go%0A%20%20%20%20from%20plotly.subplots%20import%20make_subplots%0A%20%20%20%20import%20torch%0A%20%20%20%20import%20torch.nn%20as%20nn%0A%0A%20%20%20%20return%20go%2C%20make_subplots%2C%20mo%2C%20nn%2C%20np%2C%20pd%2C%20time%2C%20torch%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%5B%E2%86%90%2052%20Multi-Head%20Attention%5D(52_multi_head_attention.py)%20%7C%20%5BIndex%5D(..%2Findex.html)%20%7C%20%5B54%20Decoding%20Strategies%20%E2%86%92%5D(54_decoding_strategies.py)%0A%0A%20%20%20%20%23%2053.%20Layer%20Normalization%20vs%20RMSNorm%3A%20Internal%20Covariate%20Shift%20and%20Modern%20Activation%20Scaling%0A%0A%20%20%20%20%23%23%23%20Executive%20Summary%0A%0A%20%20%20%20In%20deep%20neural%20architectures%2C%20continuous%20parameter%20updates%20during%20optimization%20cause%20the%20distribution%20of%20each%20layer's%20inputs%20to%20drift%20throughout%20training%E2%80%94a%20phenomenon%20known%20as%20internal%20covariate%20shift.%20While%20**Batch%20Normalization**%20(Ioffe%20%26%20Szegedy%2C%202015)%20mitigated%20this%20in%20convolutional%20networks%20by%20computing%20mini-batch%20statistics%2C%20it%20critically%20degrades%20in%20sequential%20modeling%20due%20to%20dynamic%20token%20lengths%20and%20inference%20batch-size%20variability.%0A%0A%20%20%20%20**Layer%20Normalization%20(LayerNorm)**%20(Ba%2C%20Kiros%2C%20%26%20Hinton%2C%202016)%20resolved%20this%20by%20computing%20mean%20and%20variance%20across%20the%20feature%20dimension%20for%20each%20individual%20token.%20However%2C%20modern%20frontier%20architectures%20(LLaMA-3%2C%20Mistral%2C%20Gemma%2C%20DeepSeek)%20have%20universally%20replaced%20LayerNorm%20with%20**Root%20Mean%20Square%20Normalization%20(RMSNorm)**%20(Zhang%20%26%20Sennrich%2C%202019).%20RMSNorm%20demonstrates%20that%20the%20training%20stabilization%20of%20LayerNorm%20arises%20almost%20entirely%20from%20scaling%20invariance%20rather%20than%20mean-centering%2C%20enabling%20a%20computationally%20streamlined%20formulation%20that%20cuts%20memory%20operations%20and%20improves%20training%20throughput%20by%20up%20to%20%2450%5C%25%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20%5Bb%5D%20Mathematical%20Foundations%20and%20Invariance%20Derivations%0A%0A%20%20%20%20%23%23%23%201.%20Classical%20Layer%20Normalization%20(LayerNorm)%0A%0A%20%20%20%20Let%20%24x%20%5Cin%20%5Cmathbb%7BR%7D%5Ed%24%20represent%20an%20activation%20vector%20for%20a%20single%20token%20at%20a%20given%20layer.%20LayerNorm%20normalizes%20%24x%24%20across%20its%20%24d%24%20hidden%20dimensions%20using%20both%20the%20sample%20mean%20%24%5Cmu%24%20and%20sample%20variance%20%24%5Csigma%5E2%24%3A%0A%0A%20%20%20%20%24%24%5Cmu%20%3D%20%5Cfrac%7B1%7D%7Bd%7D%20%5Csum_%7Bi%3D1%7D%5Ed%20x_i%2C%20%5Cqquad%20%5Csigma%5E2%20%3D%20%5Cfrac%7B1%7D%7Bd%7D%20%5Csum_%7Bi%3D1%7D%5Ed%20(x_i%20-%20%5Cmu)%5E2%24%24%0A%0A%20%20%20%20The%20normalized%20representation%20%24%5Chat%7Bx%7D%24%20and%20final%20affine-transformed%20output%20%24y%20%5Cin%20%5Cmathbb%7BR%7D%5Ed%24%20are%3A%0A%0A%20%20%20%20%24%24%5Chat%7Bx%7D_i%20%3D%20%5Cfrac%7Bx_i%20-%20%5Cmu%7D%7B%5Csqrt%7B%5Csigma%5E2%20%2B%20%5Cepsilon%7D%7D%2C%20%5Cqquad%20y_i%20%3D%20%5Cgamma_i%20%5Chat%7Bx%7D_i%20%2B%20%5Cbeta_i%24%24%0A%0A%20%20%20%20where%20%24%5Cepsilon%20%3E%200%24%20is%20a%20small%20numerical%20stabilization%20constant%20(typically%20%2410%5E%7B-5%7D%24%20or%20%2410%5E%7B-6%7D%24)%2C%20%24%5Cgamma%20%5Cin%20%5Cmathbb%7BR%7D%5Ed%24%20is%20a%20learnable%20gain%20vector%20initialized%20to%20%24%5Cmathbf%7B1%7D%24%2C%20and%20%24%5Cbeta%20%5Cin%20%5Cmathbb%7BR%7D%5Ed%24%20is%20a%20learnable%20bias%20vector%20initialized%20to%20%24%5Cmathbf%7B0%7D%24.%0A%0A%20%20%20%20%23%23%23%23%20Computational%20Complexity%20of%20LayerNorm%3A%0A%20%20%20%20LayerNorm%20requires%20**two%20sequential%20reduction%20passes**%20across%20the%20%24d%24-dimensional%20feature%20vector%3A%0A%20%20%20%201.%20Pass%201%3A%20Sum%20elements%20to%20calculate%20%24%5Cmu%20%3D%20%5Cfrac%7B1%7D%7Bd%7D%5Csum%20x_i%24.%0A%20%20%20%202.%20Pass%202%3A%20Accumulate%20squared%20differences%20%24(x_i%20-%20%5Cmu)%5E2%24%20to%20calculate%20%24%5Csigma%5E2%24.%0A%20%20%20%20On%20modern%20GPUs%2C%20this%20introduces%20global%20memory%20synchronization%20overhead%20that%20throttles%20memory%20bandwidth.%0A%0A%20%20%20%20%23%23%23%202.%20Root%20Mean%20Square%20Normalization%20(RMSNorm)%0A%0A%20%20%20%20Zhang%20%26%20Sennrich%20(2019)%20hypothesized%20that%20the%20regularizing%20power%20of%20LayerNorm%20stems%20from%20its%20**scale%20invariance**%20property%20rather%20than%20its%20shift%20invariance%20(mean-centering).%20By%20discarding%20the%20mean-centering%20step%20and%20the%20bias%20parameter%20%24%5Cbeta%24%2C%20RMSNorm%20scales%20activations%20strictly%20by%20their%20root-mean-square%20magnitude%3A%0A%0A%20%20%20%20%24%24%5Coperatorname%7BRMS%7D(x)%20%3D%20%5Csqrt%7B%5Cfrac%7B1%7D%7Bd%7D%20%5Csum_%7Bi%3D1%7D%5Ed%20x_i%5E2%20%2B%20%5Cepsilon%7D%24%24%0A%0A%20%20%20%20The%20normalized%20vector%20%24%5Cbar%7Bx%7D%24%20and%20scaled%20output%20%24y%24%20are%3A%0A%0A%20%20%20%20%24%24%5Cbar%7Bx%7D_i%20%3D%20%5Cfrac%7Bx_i%7D%7B%5Coperatorname%7BRMS%7D(x)%7D%2C%20%5Cqquad%20y_i%20%3D%20%5Cgamma_i%20%5Cbar%7Bx%7D_i%24%24%0A%0A%20%20%20%20where%20%24%5Cgamma%20%5Cin%20%5Cmathbb%7BR%7D%5Ed%24%20is%20the%20sole%20learnable%20gain%20parameter.%0A%0A%20%20%20%20%23%23%23%23%20Architectural%20Advantages%20of%20RMSNorm%3A%0A%20%20%20%201.%20**Single-Pass%20Reduction**%3A%20Computing%20%24%5Coperatorname%7BRMS%7D(x)%24%20requires%20only%20a%20single%20pass%20%24%5Csum_%7Bi%3D1%7D%5Ed%20x_i%5E2%24%2C%20completely%20bypassing%20the%20mean%20calculation%20and%20eliminating%20intermediate%20GPU%20synchronizations.%0A%20%20%20%202.%20**Reduced%20Parameter%20Count**%3A%20By%20eliminating%20the%20bias%20vector%20%24%5Cbeta%20%5Cin%20%5Cmathbb%7BR%7D%5Ed%24%2C%20RMSNorm%20saves%20%24d%24%20parameters%20per%20normalization%20layer.%20Across%20an%2080-layer%20model%20with%20%24d%20%3D%208192%24%20(two%20norms%20per%20Transformer%20block)%2C%20this%20saves%20over%20%241.3%24%20million%20parameters%20with%20zero%20expressive%20penalty.%0A%20%20%20%203.%20**Fused%20Kernel%20Speedup**%3A%20Fused%20CUDA%20%2F%20Triton%20kernels%20execute%20RMSNorm%20%2410%5C%25%24%20to%20%2450%5C%25%24%20faster%20than%20LayerNorm%2C%20yielding%20significant%20wall-clock%20speedups%20over%20trillions%20of%20training%20tokens.%0A%0A%20%20%20%20%23%23%23%203.%20Invariance%20Analysis%3A%20Scale%20vs%20Shift%0A%0A%20%20%20%20A%20normalization%20operator%20%24%5Cmathcal%7BT%7D(x)%24%20is%20characterized%20by%20its%20mathematical%20invariance%20properties%3A%0A%0A%20%20%20%20%23%23%23%23%20Scale%20Invariance%20(%24%5Cmathcal%7BT%7D(%5Calpha%20x)%20%3D%20%5Cmathcal%7BT%7D(x)%24%20for%20%24%5Calpha%20%3E%200%24)%3A%0A%20%20%20%20-%20**LayerNorm**%3A%0A%20%20%20%20%20%20%24%24%5Cmu(%5Calpha%20x)%20%3D%20%5Calpha%20%5Cmu(x)%2C%20%5Cquad%20%5Csigma(%5Calpha%20x)%20%3D%20%5Calpha%20%5Csigma(x)%20%5Cimplies%20%5Cfrac%7B%5Calpha%20x_i%20-%20%5Calpha%20%5Cmu(x)%7D%7B%5Calpha%20%5Csigma(x)%7D%20%3D%20%5Cfrac%7Bx_i%20-%20%5Cmu(x)%7D%7B%5Csigma(x)%7D%20%3D%20%5Chat%7Bx%7D_i%24%24%0A%20%20%20%20-%20**RMSNorm**%3A%0A%20%20%20%20%20%20%24%24%5Coperatorname%7BRMS%7D(%5Calpha%20x)%20%3D%20%5Csqrt%7B%5Cfrac%7B1%7D%7Bd%7D%5Csum%20(%5Calpha%20x_i)%5E2%7D%20%3D%20%5Calpha%20%5Coperatorname%7BRMS%7D(x)%20%5Cimplies%20%5Cfrac%7B%5Calpha%20x_i%7D%7B%5Calpha%20%5Coperatorname%7BRMS%7D(x)%7D%20%3D%20%5Cfrac%7Bx_i%7D%7B%5Coperatorname%7BRMS%7D(x)%7D%20%3D%20%5Cbar%7Bx%7D_i%24%24%0A%20%20%20%20Both%20LayerNorm%20and%20RMSNorm%20exhibit%20strict%20scale%20invariance.%20If%20forward%20activations%20are%20scaled%20by%20%24%5Calpha%24%2C%20the%20normalized%20activations%20remain%20invariant%2C%20bounding%20forward%20signal%20propagation%20and%20stabilizing%20backpropagated%20gradients.%0A%0A%20%20%20%20%23%23%23%23%20Shift%20Invariance%20(%24%5Cmathcal%7BT%7D(x%20%2B%20c)%20%3D%20%5Cmathcal%7BT%7D(x)%24%20for%20%24c%20%5Cin%20%5Cmathbb%7BR%7D%24)%3A%0A%20%20%20%20-%20**LayerNorm**%3A%0A%20%20%20%20%20%20%24%24%5Cmu(x%20%2B%20c)%20%3D%20%5Cmu(x)%20%2B%20c%20%5Cimplies%20(x_i%20%2B%20c)%20-%20(%5Cmu(x)%20%2B%20c)%20%3D%20x_i%20-%20%5Cmu(x)%24%24%0A%20%20%20%20%20%20LayerNorm%20is%20strictly%20shift-invariant.%0A%20%20%20%20-%20**RMSNorm**%3A%0A%20%20%20%20%20%20%24%24%5Coperatorname%7BRMS%7D(x%20%2B%20c)%20%5Cneq%20%5Coperatorname%7BRMS%7D(x)%24%24%0A%20%20%20%20%20%20RMSNorm%20is%20**not**%20shift-invariant.%20However%2C%20in%20deep%20transformers%2C%20representations%20are%20naturally%20zero-centered%20around%20residual%20stream%20additions%2C%20rendering%20shift%20invariance%20redundant.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(go%2C%20make_subplots%2C%20mo%2C%20np)%3A%0A%20%20%20%20%23%20Simulate%20an%20unnormalized%20activation%20vector%20across%2032%20hidden%20dimensions%0A%20%20%20%20np.random.seed(42)%0A%20%20%20%20dim_demo%20%3D%2032%0A%20%20%20%20%23%20Skewed%20activation%20with%20large%20positive%20mean%20shift%20(simulating%20uncentered%20feed-forward%20output)%0A%20%20%20%20raw_activations%20%3D%20np.random.normal(loc%3D3.5%2C%20scale%3D2.0%2C%20size%3Ddim_demo)%0A%20%20%20%20raw_activations%5B5%5D%20%3D%2012.0%20%20%23%20Outlier%20spike%0A%20%20%20%20raw_activations%5B18%5D%20%3D%20-4.0%0A%0A%20%20%20%20%23%201.%20Pure%20NumPy%20LayerNorm%0A%20%20%20%20ln_mean%20%3D%20np.mean(raw_activations)%0A%20%20%20%20ln_var%20%3D%20np.var(raw_activations)%0A%20%20%20%20ln_norm%20%3D%20(raw_activations%20-%20ln_mean)%20%2F%20np.sqrt(ln_var%20%2B%201e-6)%0A%0A%20%20%20%20%23%202.%20Pure%20NumPy%20RMSNorm%0A%20%20%20%20rms_val%20%3D%20np.sqrt(np.mean(raw_activations**2)%20%2B%201e-6)%0A%20%20%20%20rms_norm%20%3D%20raw_activations%20%2F%20rms_val%0A%0A%20%20%20%20%23%20Theoretical%20runtime%20comparison%20across%20dimension%20sizes%20(simulated%20FLOPs%20%2F%20memory%20reads)%0A%20%20%20%20dim_bench%20%3D%20%5B256%2C%20512%2C%201024%2C%202048%2C%204096%2C%208192%5D%0A%20%20%20%20%23%20Relative%20memory%20access%20passes%20(LayerNorm%20%3D%202%20passes%20%2B%202%20params%2C%20RMSNorm%20%3D%201%20pass%20%2B%201%20param)%0A%20%20%20%20ln_ops%20%3D%20%5Bd%20*%204%20for%20d%20in%20dim_bench%5D%20%20%23%202%20reads%2C%201%20write%2C%201%20affine%0A%20%20%20%20rms_ops%20%3D%20%5Bd%20*%202.5%20for%20d%20in%20dim_bench%5D%20%20%23%201%20read%2C%201%20write%2C%200.5%20affine%0A%0A%20%20%20%20fig%20%3D%20make_subplots(%0A%20%20%20%20%20%20%20%20rows%3D1%2C%0A%20%20%20%20%20%20%20%20cols%3D2%2C%0A%20%20%20%20%20%20%20%20subplot_titles%3D%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%3Cb%3EActivation%20Profiles%3A%20Raw%20vs%20LayerNorm%20vs%20RMSNorm%3C%2Fb%3E%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%3Cb%3ETheoretical%20Memory%20Access%20Passes%20vs%20Hidden%20Dimension%3C%2Fb%3E%22%2C%0A%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20horizontal_spacing%3D0.14%2C%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20Panel%201%3A%20Activation%20Comparison%0A%20%20%20%20x_axis%20%3D%20np.arange(1%2C%20dim_demo%20%2B%201)%0A%20%20%20%20fig.add_trace(%0A%20%20%20%20%20%20%20%20go.Scatter(%0A%20%20%20%20%20%20%20%20%20%20%20%20x%3Dx_axis%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20y%3Draw_activations%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mode%3D%22markers%2Blines%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22%239CA3AF%22%2C%20dash%3D%22dot%22%2C%20width%3D1.5)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20marker%3Ddict(size%3D6)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22Raw%20Input%20(Mean%3D3.5)%22%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D1%2C%0A%20%20%20%20)%0A%20%20%20%20fig.add_trace(%0A%20%20%20%20%20%20%20%20go.Scatter(%0A%20%20%20%20%20%20%20%20%20%20%20%20x%3Dx_axis%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20y%3Dln_norm%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mode%3D%22markers%2Blines%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22%231D4ED8%22%2C%20width%3D2.5)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20marker%3Ddict(size%3D7)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22LayerNorm%20(Mean%3D0%2C%20Std%3D1)%22%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D1%2C%0A%20%20%20%20)%0A%20%20%20%20fig.add_trace(%0A%20%20%20%20%20%20%20%20go.Scatter(%0A%20%20%20%20%20%20%20%20%20%20%20%20x%3Dx_axis%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20y%3Drms_norm%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mode%3D%22markers%2Blines%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22%230D9488%22%2C%20width%3D2.5)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20marker%3Ddict(size%3D7)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22RMSNorm%20(RMS%3D1.0)%22%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D1%2C%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20Zero%20line%20on%20panel%201%0A%20%20%20%20fig.add_hline(y%3D0.0%2C%20line%3Ddict(color%3D%22%236B7280%22%2C%20width%3D1%2C%20dash%3D%22dash%22)%2C%20row%3D1%2C%20col%3D1)%0A%0A%20%20%20%20%23%20Panel%202%3A%20Computational%20%2F%20Memory%20Access%20Overhead%0A%20%20%20%20fig.add_trace(%0A%20%20%20%20%20%20%20%20go.Bar(%0A%20%20%20%20%20%20%20%20%20%20%20%20x%3D%5Bstr(d)%20for%20d%20in%20dim_bench%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20y%3Dln_ops%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22LayerNorm%20Memory%20Ops%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20marker_color%3D%22%231D4ED8%22%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D2%2C%0A%20%20%20%20)%0A%20%20%20%20fig.add_trace(%0A%20%20%20%20%20%20%20%20go.Bar(%0A%20%20%20%20%20%20%20%20%20%20%20%20x%3D%5Bstr(d)%20for%20d%20in%20dim_bench%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20y%3Drms_ops%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22RMSNorm%20Memory%20Ops%20(37.5%25%20Reduction)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20marker_color%3D%22%230D9488%22%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%0A%20%20%20%20%20%20%20%20col%3D2%2C%0A%20%20%20%20)%0A%0A%20%20%20%20fig.update_xaxes(title_text%3D%22Hidden%20Feature%20Dimension%20Index%22%2C%20row%3D1%2C%20col%3D1)%0A%20%20%20%20fig.update_yaxes(title_text%3D%22Activation%20Value%22%2C%20row%3D1%2C%20col%3D1)%0A%20%20%20%20fig.update_xaxes(title_text%3D%22Hidden%20Dimension%20d_model%22%2C%20row%3D1%2C%20col%3D2)%0A%20%20%20%20fig.update_yaxes(title_text%3D%22Relative%20Memory%20Operations%20(kOps)%22%2C%20row%3D1%2C%20col%3D2)%0A%0A%20%20%20%20fig.update_layout(%0A%20%20%20%20%20%20%20%20template%3D%22plotly_white%22%2C%0A%20%20%20%20%20%20%20%20height%3D500%2C%0A%20%20%20%20%20%20%20%20margin%3Ddict(l%3D40%2C%20r%3D40%2C%20t%3D70%2C%20b%3D50)%2C%0A%20%20%20%20%20%20%20%20legend%3Ddict(orientation%3D%22h%22%2C%20yanchor%3D%22bottom%22%2C%20y%3D-0.28%2C%20xanchor%3D%22center%22%2C%20x%3D0.5)%2C%0A%20%20%20%20%20%20%20%20barmode%3D%22group%22%2C%0A%20%20%20%20)%0A%0A%20%20%20%20viz%20%3D%20mo.ui.plotly(fig)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20nn%2C%20np%2C%20pd%2C%20time%2C%20torch)%3A%0A%20%20%20%20%23%20Vectorized%20NumPy%20Implementations%0A%20%20%20%20def%20numpy_layernorm(x%2C%20gamma%3DNone%2C%20beta%3DNone%2C%20eps%3D1e-6)%3A%0A%20%20%20%20%20%20%20%20mean%20%3D%20np.mean(x%2C%20axis%3D-1%2C%20keepdims%3DTrue)%0A%20%20%20%20%20%20%20%20var%20%3D%20np.var(x%2C%20axis%3D-1%2C%20keepdims%3DTrue)%0A%20%20%20%20%20%20%20%20x_hat%20%3D%20(x%20-%20mean)%20%2F%20np.sqrt(var%20%2B%20eps)%0A%20%20%20%20%20%20%20%20if%20gamma%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20x_hat%20%3D%20x_hat%20*%20gamma%0A%20%20%20%20%20%20%20%20if%20beta%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20x_hat%20%3D%20x_hat%20%2B%20beta%0A%20%20%20%20%20%20%20%20return%20x_hat%0A%0A%20%20%20%20def%20numpy_rmsnorm(x%2C%20gamma%3DNone%2C%20eps%3D1e-6)%3A%0A%20%20%20%20%20%20%20%20rms%20%3D%20np.sqrt(np.mean(x**2%2C%20axis%3D-1%2C%20keepdims%3DTrue)%20%2B%20eps)%0A%20%20%20%20%20%20%20%20x_hat%20%3D%20x%20%2F%20rms%0A%20%20%20%20%20%20%20%20if%20gamma%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20x_hat%20%3D%20x_hat%20*%20gamma%0A%20%20%20%20%20%20%20%20return%20x_hat%0A%0A%20%20%20%20%23%20Verification%201%3A%20Scale%20Invariance%20and%20Shift%20Invariance%20Axioms%0A%20%20%20%20np.random.seed(1337)%0A%20%20%20%20_x%20%3D%20np.random.normal(2.0%2C%203.0%2C%20size%3D(10%2C%2064))%0A%20%20%20%20_scale_factor%20%3D%2042.5%0A%20%20%20%20_shift_constant%20%3D%2017.8%0A%0A%20%20%20%20%23%20Scale%20invariance%20test%0A%20%20%20%20_ln_base%20%3D%20numpy_layernorm(_x)%0A%20%20%20%20_ln_scaled%20%3D%20numpy_layernorm(_x%20*%20_scale_factor)%0A%20%20%20%20_ln_scale_diff%20%3D%20np.max(np.abs(_ln_base%20-%20_ln_scaled))%0A%0A%20%20%20%20_rms_base%20%3D%20numpy_rmsnorm(_x)%0A%20%20%20%20_rms_scaled%20%3D%20numpy_rmsnorm(_x%20*%20_scale_factor)%0A%20%20%20%20_rms_scale_diff%20%3D%20np.max(np.abs(_rms_base%20-%20_rms_scaled))%0A%0A%20%20%20%20%23%20Shift%20invariance%20test%0A%20%20%20%20_ln_shifted%20%3D%20numpy_layernorm(_x%20%2B%20_shift_constant)%0A%20%20%20%20_ln_shift_diff%20%3D%20np.max(np.abs(_ln_base%20-%20_ln_shifted))%0A%0A%20%20%20%20_rms_shifted%20%3D%20numpy_rmsnorm(_x%20%2B%20_shift_constant)%0A%20%20%20%20_rms_shift_diff%20%3D%20np.max(np.abs(_rms_base%20-%20_rms_shifted))%0A%0A%20%20%20%20df_invariance%20%3D%20pd.DataFrame(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Mathematical_Property%22%3A%20%22Scale%20Invariance%3A%20Norm(alpha%20*%20x)%20%3D%3D%20Norm(x)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22LayerNorm_Max_Diff%22%3A%20f%22%7B_ln_scale_diff%3A.8e%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22RMSNorm_Max_Diff%22%3A%20f%22%7B_rms_scale_diff%3A.8e%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Theoretical_Status%22%3A%20%22Both%20Strictly%20Scale%20Invariant%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Mathematical_Property%22%3A%20%22Shift%20Invariance%3A%20Norm(x%20%2B%20c)%20%3D%3D%20Norm(x)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22LayerNorm_Max_Diff%22%3A%20f%22%7B_ln_shift_diff%3A.8e%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22RMSNorm_Max_Diff%22%3A%20f%22%7B_rms_shift_diff%3A.4f%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Theoretical_Status%22%3A%20%22LayerNorm%20Invariant%20%7C%20RMSNorm%20Non-Invariant%20(By%20Design)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Mathematical_Property%22%3A%20%22Mean-Centering%20Property%3A%20Mean(Norm(x))%20%3D%3D%200%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22LayerNorm_Max_Diff%22%3A%20f%22%7Bnp.max(np.abs(np.mean(_ln_base%2C%20axis%3D-1)))%3A.8e%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22RMSNorm_Max_Diff%22%3A%20f%22%7Bnp.max(np.abs(np.mean(_rms_base%2C%20axis%3D-1)))%3A.4f%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Theoretical_Status%22%3A%20%22LayerNorm%20Exactly%200%20%7C%20RMSNorm%20Inherits%20Scaled%20Mean%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20Verification%202%3A%20PyTorch%20Modules%20and%20Gradient%20Verification%0A%20%20%20%20class%20PyTorchRMSNorm(nn.Module)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20dim%2C%20eps%3D1e-6)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20super().__init__()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.eps%20%3D%20eps%0A%20%20%20%20%20%20%20%20%20%20%20%20self.weight%20%3D%20nn.Parameter(torch.ones(dim))%0A%0A%20%20%20%20%20%20%20%20def%20forward(self%2C%20x)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20norm_x%20%3D%20torch.mean(x%20*%20x%2C%20dim%3D-1%2C%20keepdim%3DTrue)%0A%20%20%20%20%20%20%20%20%20%20%20%20x_normed%20%3D%20x%20*%20torch.rsqrt(norm_x%20%2B%20self.eps)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.weight%20*%20x_normed%0A%0A%20%20%20%20torch.manual_seed(42)%0A%20%20%20%20_d%20%3D%20128%0A%20%20%20%20pt_ln%20%3D%20nn.LayerNorm(_d)%0A%20%20%20%20pt_rms%20%3D%20PyTorchRMSNorm(_d)%0A%0A%20%20%20%20pt_x%20%3D%20torch.randn(4%2C%2016%2C%20_d%2C%20requires_grad%3DTrue)%0A%0A%20%20%20%20%23%20Run%20PyTorch%20modules%0A%20%20%20%20out_ln%20%3D%20pt_ln(pt_x)%0A%20%20%20%20loss_ln%20%3D%20out_ln.sum()%0A%20%20%20%20loss_ln.backward()%0A%20%20%20%20grad_ln_norm%20%3D%20pt_x.grad.norm().item()%0A%0A%20%20%20%20pt_x.grad.zero_()%0A%20%20%20%20out_rms%20%3D%20pt_rms(pt_x)%0A%20%20%20%20loss_rms%20%3D%20out_rms.sum()%0A%20%20%20%20loss_rms.backward()%0A%20%20%20%20grad_rms_norm%20%3D%20pt_x.grad.norm().item()%0A%0A%20%20%20%20df_modules%20%3D%20pd.DataFrame(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Normalization_Layer%22%3A%20%22PyTorch%20LayerNorm%20(nn.LayerNorm)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Parameters_per_Layer%22%3A%20f%22%7B2%20*%20_d%7D%20(gamma%20%2B%20beta)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Requires_Bias_Vector%22%3A%20%22Yes%20(beta%20in%20R%5Ed)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Backprop_Grad_Norm%22%3A%20f%22%7Bgrad_ln_norm%3A.4f%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Reduction_Passes%22%3A%20%222%20passes%20(Mean%20and%20Variance)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Normalization_Layer%22%3A%20%22PyTorch%20RMSNorm%20(Modern%20LLM)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Parameters_per_Layer%22%3A%20f%22%7B_d%7D%20(gamma%20only)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Requires_Bias_Vector%22%3A%20%22No%20(Zero%20Bias)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Backprop_Grad_Norm%22%3A%20f%22%7Bgrad_rms_norm%3A.4f%7D%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Reduction_Passes%22%3A%20%221%20pass%20(Mean%20Square)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20Verification%203%3A%20Empirical%20Forward-Pass%20Latency%20Benchmark%0A%20%20%20%20%23%20Simulate%20LLM%20tensor%3A%20Batch%3D8%2C%20Sequence%20Length%3D128%2C%20Hidden%20Dim%3D1024%0A%20%20%20%20_b_test%2C%20_s_test%2C%20_d_test%20%3D%208%2C%20128%2C%201024%0A%20%20%20%20bench_x%20%3D%20torch.randn(_b_test%2C%20_s_test%2C%20_d_test)%0A%20%20%20%20bench_ln%20%3D%20nn.LayerNorm(_d_test)%0A%20%20%20%20bench_rms%20%3D%20PyTorchRMSNorm(_d_test)%0A%0A%20%20%20%20%23%20Warmup%0A%20%20%20%20for%20_%20in%20range(50)%3A%0A%20%20%20%20%20%20%20%20_%20%3D%20bench_ln(bench_x)%0A%20%20%20%20%20%20%20%20_%20%3D%20bench_rms(bench_x)%0A%0A%20%20%20%20%23%20Benchmark%20500%20iterations%0A%20%20%20%20n_iters%20%3D%20500%0A%20%20%20%20t0%20%3D%20time.perf_counter()%0A%20%20%20%20for%20_%20in%20range(n_iters)%3A%0A%20%20%20%20%20%20%20%20_%20%3D%20bench_ln(bench_x)%0A%20%20%20%20t_ln%20%3D%20(time.perf_counter()%20-%20t0)%20*%201000.0%0A%0A%20%20%20%20t0%20%3D%20time.perf_counter()%0A%20%20%20%20for%20_%20in%20range(n_iters)%3A%0A%20%20%20%20%20%20%20%20_%20%3D%20bench_rms(bench_x)%0A%20%20%20%20t_rms%20%3D%20(time.perf_counter()%20-%20t0)%20*%201000.0%0A%0A%20%20%20%20speedup%20%3D%20(t_ln%20-%20t_rms)%20%2F%20t_ln%20*%20100.0%0A%0A%20%20%20%20df_benchmark%20%3D%20pd.DataFrame(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Normalization_Type%22%3A%20%22Standard%20LayerNorm%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Tensor_Dimension%22%3A%20f%22(%7B_b_test%7D%2C%20%7B_s_test%7D%2C%20%7B_d_test%7D)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Total_500_Iter_Time%22%3A%20f%22%7Bt_ln%3A.2f%7D%20ms%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Mean_Latency_per_Pass%22%3A%20f%22%7Bt_ln%20%2F%20n_iters%20*%201000%3A.2f%7D%20us%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Speedup_vs_LayerNorm%22%3A%20%22Baseline%20(1.0x)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Normalization_Type%22%3A%20%22Streamlined%20RMSNorm%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Tensor_Dimension%22%3A%20f%22(%7B_b_test%7D%2C%20%7B_s_test%7D%2C%20%7B_d_test%7D)%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Total_500_Iter_Time%22%3A%20f%22%7Bt_rms%3A.2f%7D%20ms%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Mean_Latency_per_Pass%22%3A%20f%22%7Bt_rms%20%2F%20n_iters%20*%201000%3A.2f%7D%20us%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Speedup_vs_LayerNorm%22%3A%20f%22%7Bspeedup%3A.1f%7D%25%20faster%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%7D%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20table_invariance%20%3D%20mo.ui.table(df_invariance)%0A%20%20%20%20table_modules%20%3D%20mo.ui.table(df_modules)%0A%20%20%20%20table_bench%20%3D%20mo.ui.table(df_benchmark)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A
bb351e66f81142e0839bf1de7733a6b1