gsnn.gsnn.models.GroupRMSNorm

class gsnn.gsnn.models.GroupRMSNorm(*args: Any, **kwargs: Any)[source]

Bases: Module

Applies Root Mean Square normalization within each channel group.

RMSNorm normalizes using only the RMS (root mean square) without mean centering, making it simpler and more stable than layer normalization, especially for small batch sizes.

Parameters:

channel_groups (list or tensor) – Specifies which group each channel belongs to. For example, [0,0,1,1,2,2] specifies 3 groups with 2 channels each.
eps (float) – Small value to avoid division by zero. Default: 1e-6
affine (bool) – If True, applies learnable scale parameter. Default: True

__init__(channel_groups, eps=1e-06, affine=True)[source]

Methods

`__init__`(channel_groups[, eps, affine])
`forward`(x)	Normalize by group RMS; accepts `(B, C)` or `(B, C, 1)`.

forward(x)[source]: Normalize by group RMS; accepts (B, C) or (B, C, 1).