Struct RMSNorm 

pub struct RMSNorm { /* private fields */ }

Root Mean Square Layer Normalization (Zhang & Sennrich, 2019).

A simplified version of LayerNorm that only uses the root mean square for normalization, without centering (no mean subtraction). This is faster than LayerNorm while achieving similar results.

y = x / RMS(x) * gamma
RMS(x) = sqrt(mean(x^2) + eps)

Used in LLaMA, Gemma, and other modern transformers.

Example

use aprender::nn::{RMSNorm, Module};
use aprender::autograd::Tensor;

let norm = RMSNorm::new(&[256]);  // Normalize over 256 features
let x = Tensor::randn(&[32, 10, 256]);  // [batch, seq, features]
let y = norm.forward(&x);  // Normalized

References

• Zhang, B., & Sennrich, R. (2019). Root Mean Square Layer Normalization. NeurIPS.

Implementations

impl RMSNorm

pub fn new(normalized_shape: &[usize]) -> Self

Create a new RMSNorm layer.

Arguments

• normalized_shape - Shape of the dimensions to normalize over

pub fn with_eps(normalized_shape: &[usize], eps: f32) -> Self

Create RMSNorm with custom epsilon.

pub fn without_affine(normalized_shape: &[usize]) -> Self

Create RMSNorm without learnable parameters.

pub fn normalized_shape(&self) -> &[usize]

Get the normalized shape.

pub fn eps(&self) -> f32

Get the epsilon value.

pub fn set_weight(&mut self, weight: Tensor)

Set weight tensor from external data.

Used for loading pre-trained weights.

pub fn weight(&self) -> &Tensor

Get reference to weight tensor.

pub fn placeholder(normalized_shape: &[usize]) -> Self

Create a placeholder RMSNorm layer with minimal memory allocation.

Used for lazy initialization when loading pre-trained weights. The placeholder uses 1-element tensors instead of full vectors, reducing memory from O(n) to O(1).

IMPORTANT: This layer will NOT work for inference until set_weight() is called with real weights.

Trait Implementations

impl Debug for RMSNorm

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Module for RMSNorm

fn forward(&self, input: &Tensor) -> Tensor

Perform forward computation.

fn parameters(&self) -> Vec<&Tensor>

Get references to all learnable parameters.

fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Get mutable references to all learnable parameters.

fn refresh_caches(&mut self)

Refresh any cached computations after parameters have been modified.

fn train(&mut self)

Set the module to training mode.

fn eval(&mut self)

Set the module to evaluation mode.

fn training(&self) -> bool

Check if the module is in training mode.

fn zero_grad(&mut self)

Zero out gradients for all parameters.

fn num_parameters(&self) -> usize

Get the number of learnable parameters.