| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """Implementation of custom initializers for NN parameters.""" |
|
|
| from typing import Any, Optional, Sequence |
|
|
| import gin |
| import haiku as hk |
| import jax |
| import numpy as np |
|
|
|
|
| |
| Constant = gin.external_configurable(hk.initializers.Constant) |
| VarianceScaling = gin.external_configurable(hk.initializers.VarianceScaling) |
| Orthogonal = gin.external_configurable(hk.initializers.Orthogonal) |
|
|
|
|
| def _compute_fans( |
| shape: Sequence[int], |
| fan_in_axes: Optional[Sequence[int]] = None, |
| ) -> tuple[int, int]: |
| """Computes the number of input and output units for a weight shape.""" |
| |
| if len(shape) < 1: |
| fan_in = fan_out = 1 |
| elif len(shape) == 1: |
| fan_in = fan_out = shape[0] |
| elif len(shape) == 2: |
| fan_in, fan_out = shape |
| else: |
| if fan_in_axes is not None: |
| |
| fan_in = np.prod([shape[i] for i in fan_in_axes]) |
| fan_out = np.prod([s for i, s in enumerate(shape) |
| if i not in fan_in_axes]) |
| else: |
| |
| |
| receptive_field_size = np.prod(shape[:-2]) |
| fan_in = shape[-2] * receptive_field_size |
| fan_out = shape[-1] * receptive_field_size |
| return fan_in, fan_out |
|
|
|
|
| @gin.register |
| class ReducingVarianceScaling(hk.initializers.Initializer): |
| """Initializer that result in variance that reduces as width increases. |
| |
| Initializes weights that result in features with expected variance of |
| `scale / n`, where `n` corresponds to the width of the layer. This initializer |
| can be used in the output layer to achieve µ parameterization [1]. |
| |
| References: |
| [1]: https://arxiv.org/abs/2203.03466 |
| """ |
|
|
| def __init__( |
| self, |
| scale=1.0, |
| mode='fan_in', |
| distribution='truncated_normal', |
| fan_in_axes=None, |
| ): |
| """Constructs `ReducingVarianceScaling` initializer. |
| |
| Args: |
| scale: Variance scale for a width == 1 initialization. |
| mode: One of ``fan_in``, ``fan_out``, ``fan_avg`` |
| distribution: Random distribution to use. One of ``truncated_normal``, |
| ``normal`` or ``uniform``. |
| fan_in_axes: Optional sequence of int specifying which axes of the shape |
| are part of the fan-in. If none provided, then the weight is assumed |
| to be like a convolution kernel, where all leading dimensions are part |
| of the fan-in, and only the trailing dimension is part of the fan-out. |
| Useful if instantiating multi-headed attention weights. |
| """ |
| if scale < 0.0: |
| raise ValueError('`scale` must be a positive float.') |
| if mode not in {'fan_in', 'fan_out', 'fan_avg'}: |
| raise ValueError('Invalid `mode` argument:', mode) |
| distribution = distribution.lower() |
| if distribution not in {'normal', 'truncated_normal', 'uniform'}: |
| raise ValueError('Invalid `distribution` argument:', distribution) |
| self.scale = scale |
| self.mode = mode |
| self.distribution = distribution |
| self.fan_in_axes = fan_in_axes |
|
|
| def __call__(self, shape: Sequence[int], dtype: Any) -> jax.Array: |
| scale = self.scale |
| fan_in, fan_out = _compute_fans(shape, self.fan_in_axes) |
| if self.mode == 'fan_in': |
| scale /= max(1.0, fan_in) ** 2 |
| elif self.mode == 'fan_out': |
| scale /= max(1.0, fan_out) ** 2 |
| else: |
| scale /= max(1.0, (fan_in + fan_out) / 2.0) ** 2 |
|
|
| if self.distribution == 'truncated_normal': |
| stddev = np.sqrt(scale) |
| |
| |
| distribution_stddev = np.asarray(.87962566103423978, dtype=dtype) |
| stddev = stddev / distribution_stddev |
| return hk.initializers.TruncatedNormal(stddev=stddev)(shape, dtype) |
| elif self.distribution == 'normal': |
| stddev = np.sqrt(scale) |
| return hk.initializers.RandomNormal(stddev=stddev)(shape, dtype) |
| else: |
| limit = np.sqrt(3.0 * scale) |
| uniform_init = hk.initializers.RandomUniform(minval=-limit, maxval=limit) |
| return uniform_init(shape, dtype) |
|
|
