SCRNCell
LuxRecurrentLayers.SCRNCell
— TypeSCRNCell(in_dims => out_dims;
use_bias=true, train_state=false, train_memory=false,
init_bias=nothing, init_weight=nothing, init_recurrent_weight=nothing,
init_context_weight=nothing, init_state=zeros32, init_memory=zeros32)
Structurally contraint recurrent unit.
Equations
\[\begin{aligned} \mathbf{s}(t) &= (1 - \alpha) \left( \mathbf{W}_{ih}^{s} \mathbf{x}(t) + \mathbf{b}_{ih}^{s} \right) + \alpha \, \mathbf{s}(t-1), \\ \mathbf{h}(t) &= \sigma\left( \mathbf{W}_{ch}^{h} \mathbf{s}(t) + \mathbf{b}_{ch}^{h} + \mathbf{W}_{ih}^{h} \mathbf{x}(t) + \mathbf{b}_{ih}^{h} + \mathbf{W}_{hh}^{h} \mathbf{h}(t-1) + \mathbf{b}_{hh}^{h} \right), \\ \mathbf{y}(t) &= f\left( \mathbf{W}_{ch}^{y} \mathbf{s}(t) + \mathbf{b}_{ch}^{y} + \mathbf{W}_{hh}^{y} \mathbf{h}(t) + \mathbf{b}_{hh}^{y} \right) \end{aligned}\]
Arguments
in_dims
: Input Dimensionout_dims
: Output (Hidden State & Memory) Dimension
Keyword Arguments
use_bias
: Flag to use bias in the computation. Default set totrue
.train_state
: Flag to set the initial hidden state as trainable. Default set tofalse
.train_memory
: Flag to set the initial memory state as trainable. Default set tofalse
.init_bias
: Initializer for input-to-hidden biases $\mathbf{b}_{ih}^{s}, \mathbf{b}_{ih}^{h}$. Must be a tuple containing 2 functions. If a single value is passed, it is copied into a 2-element tuple. If set tonothing
, biases are initialized from a uniform distribution within[-bound, bound]
, wherebound = inv(sqrt(out_dims))
. The functions are applied in order: the first initializes $\mathbf{b}_{ih}^{s}$, the second $\mathbf{b}_{ih}^{h}$. Default set tonothing
.init_recurrent_bias
: Initializer for hidden-to-hidden biases $\mathbf{b}_{hh}^{h}, \mathbf{b}_{hh}^{y}$. Must be a tuple containing 2 functions. If a single value is passed, it is copied into a 2-element tuple. If set tonothing
, biases are initialized from a uniform distribution within[-bound, bound]
wherebound = inv(sqrt(out_dims))
. Default set tonothing
.init_context_bias
: Initializer for context biases $\mathbf{b}_{ch}^{h}, \mathbf{b}_{ch}^{y}$. Must be a tuple containing 2 functions. If a single value is passed, it is copied into a 2-element tuple. If set tonothing
, biases are initialized from a uniform distribution within[-bound, bound]
wherebound = inv(sqrt(out_dims))
. Default set tonothing
.init_weight
: Initializer for input-to-hidden weights $\mathbf{W}_{ih}^{s}, \mathbf{W}_{ih}^{h}$. Must be a tuple containing 2 functions. If a single value is passed, it is copied into a 2-element tuple. If set tonothing
, weights are initialized from a uniform distribution within[-bound, bound]
, wherebound = inv(sqrt(out_dims))
. The functions are applied in order: the first initializes $\mathbf{W}_{ih}^{s}$, the second $\mathbf{W}_{ih}^{h}$. Default set tonothing
.init_recurrent_weight
: Initializer for hidden-to-hidden weights $\mathbf{W}_{hh}^{h}, \mathbf{W}_{hh}^{y}$. Must be a tuple containing 2 functions. If a single value is passed, it is copied into a 2-element tuple. If set tonothing
, weights are initialized from a uniform distribution within[-bound, bound]
wherebound = inv(sqrt(out_dims))
. Default set tonothing
.init_context_weight
: Initializer for context weights $\mathbf{W}_{ch}^{h}, \mathbf{W}_{ch}^{y}$. Must be a tuple containing 2 functions. If a single value is passed, it is copied into a 2-element tuple. If set tonothing
, weights are initialized from a uniform distribution within[-bound, bound]
wherebound = inv(sqrt(out_dims))
. Default set tonothing
.init_state
: Initializer for hidden state. Default set tozeros32
.init_memory
: Initializer for memory. Default set tozeros32
.
Inputs
- Case 1a: Only a single input
x
of shape(in_dims, batch_size)
,train_state
is set tofalse
,train_memory
is set tofalse
- Creates a hidden state usinginit_state
, hidden memory usinginit_memory
and proceeds to Case 2. - Case 1b: Only a single input
x
of shape(in_dims, batch_size)
,train_state
is set totrue
,train_memory
is set tofalse
- Repeatshidden_state
vector from the parameters to match the shape ofx
, creates hidden memory usinginit_memory
and proceeds to Case 2. - Case 1c: Only a single input
x
of shape(in_dims, batch_size)
,train_state
is set tofalse
,train_memory
is set totrue
- Creates a hidden state usinginit_state
, repeats the memory vector from parameters to match the shape ofx
and proceeds to Case 2. - Case 1d: Only a single input
x
of shape(in_dims, batch_size)
,train_state
is set totrue
,train_memory
is set totrue
- Repeats the hidden state and memory vectors from the parameters to match the shape ofx
and proceeds to Case 2. - Case 2: Tuple
(x, (h, c))
is provided, then the output and a tuple containing the updated hidden state and memory is returned.
Returns
Tuple Containing
- Output $h_{new}$ of shape
(out_dims, batch_size)
- Tuple containing new hidden state $h_{new}$ and new memory $c_{new}$
- Output $h_{new}$ of shape
Updated model state
Parameters
weight_ch
: Context-to-hidden weights $\{ \mathbf{W}_{ch}^{h}, \mathbf{W}_{ch}^{y} \}$weight_ih
: Input-to-hidden weights $\{ \mathbf{W}_{ih}^{s}, \mathbf{W}_{ih}^{h} \}$weight_hh
: Hidden-to-hidden weights $\{ \mathbf{W}_{hh}^{h}, \mathbf{W}_{hh}^{y} \}$bias_ch
: Context-to-hidden biases (not present ifuse_bias=false
) $\{ \mathbf{b}_{ch}^{h}, \mathbf{b}_{ch}^{y} \}$bias_ih
: Input-to-hidden biases (not present ifuse_bias=false
) $\{ \mathbf{b}_{ih}^{s}, \mathbf{b}_{ih}^{h} \}$bias_hh
: Hidden-to-hidden biases (not present ifuse_bias=false
) $\{ \mathbf{b}_{hh}^{h}, \mathbf{b}_{hh}^{y} \}$alpha
: Initial context strengthhidden_state
: Initial hidden state vector (not present iftrain_state=false
)memory
: Initial memory vector (not present iftrain_memory=false
)
States
rng
: Controls the randomness (if any) in the initial state generation