MUT2Cell

MUT2Cell(in_dims => out_dims;
    use_bias=true, train_state=false,
    init_bias=nothing, init_recurrent_bias=nothing,
    init_weight=nothing, init_recurrent_weight=nothing,
    init_state=zeros32)

Mutated unit 2 cell.

Equations

\[\begin{aligned} \mathbf{z}(t) &= \sigma\left( \mathbf{W}_{ih}^{z} \mathbf{x}(t) + \mathbf{b}_{ih}^{z} + \mathbf{W}_{hh}^{z} \mathbf{h}(t) + \mathbf{b}_{hh}^{z} \right), \\ \mathbf{r}(t) &= \sigma\left( \mathbf{W}_{ih}^{r} \mathbf{x}(t) + \mathbf{b}_{ih}^{r} + \mathbf{W}_{hh}^{r} \mathbf{h}(t) + \mathbf{b}_{hh}^{r} \right), \\ \mathbf{h}(t+1) &= \tanh\left( \mathbf{W}_{hh}^{h} \left( \mathbf{r}(t) \circ \mathbf{h}(t) + \mathbf{b}_{hh}^{h} \right) + \mathbf{W}_{ih}^{h} \mathbf{x}(t) + \mathbf{b}_{ih}^{h} \right) \circ \mathbf{z}(t) \\ &\quad + \mathbf{h}(t) \circ \left( 1 - \mathbf{z}(t) \right) \end{aligned}\]

Arguments

• in_dims: Input Dimension
• out_dims: Output (Hidden State & Memory) Dimension

Keyword Arguments

• use_bias: Flag to use bias in the computation. Default set to true.
• train_state: Flag to set the initial hidden state as trainable. Default set to false.
• init_bias: Initializer for input-to-hidden biases $\mathbf{b}_{ih}^{z}, \mathbf{b}_{ih}^{r}, \mathbf{b}_{ih}^{h}$. Must be a tuple containing 3 functions. If a single value is passed, it is copied into a 3-element tuple. If set to nothing, weights are initialized from a uniform distribution within [-bound, bound] where bound = inv(sqrt(out_dims)). The functions are applied in order: the first initializes $\mathbf{b}_{ih}^{z}$, the second $\mathbf{b}_{ih}^{r}$, and the third $\mathbf{b}_{ih}^{h}$.
• init_recurrent_bias: Initializer for hidden-to-hidden biases $\mathbf{b}_{hh}^{z}, \mathbf{b}_{hh}^{r}, \mathbf{b}_{hh}^{h}$. Must be a tuple containing 3 functions. If a single value is passed, it is copied into a 3-element tuple. If set to nothing, weights are initialized from a uniform distribution within [-bound, bound] where bound = inv(sqrt(out_dims)). The functions are applied in order: the first initializes $\mathbf{b}_{hh}^{z}$, the second $\mathbf{b}_{hh}^{r}$, and the third $\mathbf{b}_{hh}^{h}$.
• init_weight: Initializer for input-to-hidden weights $\mathbf{W}_{ih}^{z}, \mathbf{W}_{ih}^{r}, \mathbf{W}_{ih}^{h}$. Must be a tuple containing 3 functions. If a single value is passed, it is copied into a 3-element tuple. If set to nothing, weights are initialized from a uniform distribution within [-bound, bound] where bound = inv(sqrt(out_dims)). The functions are applied in order: the first initializes $\mathbf{W}_{ih}^{z}$, the second $\mathbf{W}_{ih}^{r}$, and the third $\mathbf{W}_{ih}^{h}$.
• init_recurrent_weight: Initializer for hidden-to-hidden weights $\mathbf{W}_{hh}^{z}, \mathbf{W}_{hh}^{r}, \mathbf{W}_{hh}^{h}$. Must be a tuple containing 3 functions. If a single value is passed, it is copied into a 3-element tuple. If set to nothing, weights are initialized from a uniform distribution within [-bound, bound] where bound = inv(sqrt(out_dims)). The functions are applied in order: the first initializes $\mathbf{W}_{hh}^{z}$, the second $\mathbf{W}_{hh}^{r}$, and the third $\mathbf{W}_{hh}^{h}$.
• init_state: Initializer for hidden state. Default set to zeros32.

Inputs

• Case 1a: Only a single input x of shape (in_dims, batch_size), train_state is set to false - Creates a hidden state using init_state and proceeds to Case 2.
• Case 1b: Only a single input x of shape (in_dims, batch_size), train_state is set to true - Repeats hidden_state from parameters to match the shape of x and proceeds to Case 2.
• Case 2: Tuple (x, (h, )) is provided, then the output and a tuple containing the updated hidden state is returned.

Returns

• Tuple containing

  • Output $h_{new}$ of shape (out_dims, batch_size)
  • Tuple containing new hidden state $h_{new}$

• Updated model state

Parameters

• weight_ih: Input-to-hidden weights $\{ \mathbf{W}_{ih}^{z}, \mathbf{W}_{ih}^{r}, \mathbf{W}_{ih}^{h} \}$
• weight_hh: Hidden-to-hidden weights $\{ \mathbf{W}_{hh}^{z}, \mathbf{W}_{hh}^{r}, \mathbf{W}_{hh}^{h} \}$
• bias_ih: Input-to-hidden biases (if use_bias=true) $\{ \mathbf{b}_{ih}^{z}, \mathbf{b}_{ih}^{r}, \mathbf{b}_{ih}^{h} \}$
• bias_hh: Hidden-to-hidden biases (if use_bias=true) $\{ \mathbf{b}_{hh}^{z}, \mathbf{b}_{hh}^{r}, \mathbf{b}_{hh}^{h} \}$
• hidden_state: Initial hidden state vector (not present if train_state=false)

States

• rng: Controls the randomness (if any) in the initial state generation