DOCSEARCHbyAlgolia
Submit Docs

RMSNorm (Root Mean Square Layer Normalization)

repository·master·Indexed 43 days ago

https://github.com/bzhanggo/rmsnorm

A lightweight alternative to Layer Normalization that improves computational efficiency in Transformers and RNNs by removing mean-centering and focusing on the Root Mean Square statistic. Provides implementations for Theano, TensorFlow, PyTorch, and Lasagne, with integration guides for Nematus, Attentive Reader, order-embedding, and CIFAR-10 classification models.

Tokens
2.3K
Snippets
3
Records
8
Agent score
5%

What's inside RMSNorm

  1. What is RMSNorm?

    master

    RMSNorm (Root Mean Square Layer Normalization) is a simplified version of LayerNorm. It achieves regularization by normalizing layer activations using the RMS statistic instead of both mean and variance.

    Key characteristics:

    • Simplification: It removes the mean-centering operation found in standard LayerNorm.
    • Mathematical Formula: $\bar{a}i = \frac{a_i}{\text{RMS}(\mathbf{a})} g_i$, where $\text{RMS}(\mathbf{a}) = \sqrt{\frac{1}{n} \sum{i=1}^{n} a_i^2}$.
    • Invariance: Unlike LayerNorm, RMSNorm is not re-centering invariant because it does not consider the mean of the inputs.
    • Efficiency: It is designed to reduce the computational overhead associated with LayerNorm, particularly in recurrent neural networks (RNNs).

  2. Integrate RMSNorm into CNN/Daily Mail Reading Comprehension

    master

    To apply RMSNorm to the bidirectional attentive reader model:

    1. Clone the Attentive Reader repository and download the data using:
      wget http://www.cs.toronto.edu/~rkiros/top4.zip
    2. In codes/att_reader/pkl_data_iterator.py, set vdir to your unzipped data directory.
    3. In codes/att_reader/layers.py, add the RMSNorm function, the rlnlstm_layer and param_init_rlnlstm functions.
    4. Register the new layer by adding 'rlnlstm': ('param_init_rlnlstm', 'rlnlstm_layer') to the layers dictionary.
    5. Train the model using the following command:
    GPUARRAY_FORCE_CUDA_DRIVER_LOAD=True THEANO_FLAGS=mode=FAST_RUN,floatX=float32,device=$device,gpuarray.preallocate=0.8 python -u train_attentive_reader.py \
    --use_dq_sims 1 --use_desc_skip_c_g 0 --dim 240 --learn_h0 1 --lr 8e-5 --truncate -1 --model "lstm_s1.npz" --batch_size 64 --optimizer "adam" --validFreq 1000 --model_dir $MDIR --use_desc_skip_c_g 1  --unit_type rlnlstm --use_bidir 1

  4. Integrate RMSNorm into Nematus (Machine Translation)

    master

    To use RMSNorm with the Nematus (v0.3) framework for machine translation, follow these steps:

    1. Clone and checkout the specific Nematus version:
      git clone https://github.com/EdinburghNLP/nematus.git
cd nematus
git checkout tags/v0.3
    2. Modify layers.py to replace the LayerNormLayer implementation with the RMSNorm version. Note that if nematus_compat is true, eps must be set to 1e-5.
    3. Train the model using the standard Nematus instructions (e.g., following wmt17_systems).

    Preprocessed datasets, training scripts, and pretrained models for RNNSearch are available at http://data.statmt.org/bzhang/neurips19_rmsnorm/.

    class LayerNormLayer(object):
    def __init__(self, layer_size, eps=1e-5):
        # TODO: If nematus_compat is true, then eps must be 1e-5!
        self.new_std = tf.get_variable('new_std', [layer_size],
                                       initializer=tf.constant_initializer(1))
        self.eps = eps
        self.layer_size = layer_size

    def forward(self, x):
        ms = tf.reduce_sum(tf.square(x), axis=-1, keep_dims=True) * 1./self.layer_size
        norm_inputs = x * tf.rsqrt(ms + self.eps)
        return norm_inputs * self.new_std

  5. Integrate RMSNorm into CIFAR-10 Classification (Lasagne/Theano)

    master

    To implement RMSNorm in a ConvPool-CNN-C architecture using Lasagne:

    1. Clone the weight_norm repository.
    2. Add the RMSNormLayer class and the rms_norm helper function to nn.py.
    3. Update train.py to handle the rms_norm type:
      elif args.norm_type=='rms_norm':
    normalizer = lambda l: nn.rms_norm(l)
    4. Download the CIFAR-10 dataset and update the paths in train.py.
    5. Train using the following command:
    GPUARRAY_FORCE_CUDA_DRIVER_LOAD=True THEANO_FLAGS=mode=FAST_RUN,floatX=float32,device=$device,gpuarray.preallocate=0.4 python train.py --norm_type rms_norm --learning_rate 0.003
    class RMSNormLayer(lasagne.layers.Layer):
    def __init__(self, incoming, b=lasagne.init.Constant(0.), g=lasagne.init.Constant(1.),
                 W=lasagne.init.Normal(0.05), nonlinearity=relu, **kwargs):
        super(RMSNormLayer, self).__init__(incoming, **kwargs)
        self.nonlinearity = nonlinearity
        k = self.input_shape[1]
        if b is not None:
            self.b = self.add_param(b, (k,), name="b", regularizable=False)
        if g is not None:
            self.g = self.add_param(g, (k,), name="g")

        if len(self.input_shape)==4:
            self.axes_to_sum = (2,3)
            self.dimshuffle_args = ['x',0,'x','x']
        else:
            self.axes_to_sum = 1
            self.dimshuffle_args = ['x',0]

    def get_output_for(self, input, **kwargs):
        meanS = T.mean(input ** 2,axis=self.axes_to_sum,keepdims=True)
        norm_input = input / T.sqrt(meanS + 1e-6)

        if hasattr(self, 'g'):
            activation = norm_input*self.g.dimshuffle(*self.dimshuffle_args)
        else:
            activation = norm_input
        if hasattr(self, 'b'):
            activation += self.b.dimshuffle(*self.dimshuffle_args)

        return self.nonlinearity(activation)

def rms_norm(layer, b=lasagne.init.Constant(0.), g=lasagne.init.Constant(1.), **kwargs):
    nonlinearity = getattr(layer, 'nonlinearity', None)
    if nonlinearity is not None:
        layer.nonlinearity = lasagne.nonlinearities.identity
    if hasattr(layer, 'b'):
        del layer.params[layer.b]
        layer.b = None
    return RMSNormLayer(layer, b, g, nonlinearity=nonlinearity, **kwargs)

  6. Integrate RMSNorm into Image-Caption Retrieval (Order-Embedding)

    master

    To use RMSNorm with the order-embedding model:

    1. Clone the order-embedding repository.
    2. Add the RMSNorm function, rlngru_layer, and param_init_rlngru to layers.py.
    3. Register the layer: add 'rlngru': ('param_init_rlngru', 'rlngru_layer') to layers.
    4. In driver.py, change the encoder from 'gru' to 'rlngru': 'encoder': 'rlngru'
    5. Follow the repository's main page instructions for training.

    Note: When evaluating using saved models, change the prefix for the rlngru model to lngru.

  7. Integrate RMSNorm into Theano-based Nematus

    master

    For Theano-based Nematus implementations, replace the layer_norm function in layers.py with the following code. Note that the bias parameter b is deletable in this implementation.

    def layer_norm(x, b, s):
    _eps = numpy_floatX(1e-5)
    norm_x = tensor.mean(x * x, axis=-1, keepdims=True)
    output = x / tensor.sqrt(norm_x + _eps)

    if x.ndim == 3:
        output = s[None, None, :] * output + b[None, None,:]
    else:
        output = s[None, :] * output + b[None,:]
    return output

  8. Access RMSNorm implementations

    master

    The repository provides different implementations depending on your framework of choice:

    • For Theano-related experiments, use layers.py.
    • For Tensorflow or PyTorch use, use rmsnorm_tensorflow/torch.py.