When implementing recurrent architectures like Long Short-Term Memory (LSTM) networks in JAX, using a standard Python for-loop to iterate over sequence time steps causes extremely long Just-In-Time (JIT) compilation times on the first run. To bypass this performance bottleneck, JAX provides the jax.lax.scan utility transformation. It accepts an initial carry state and an input array, efficiently scanning across the leading axis of the inputs to update the state and yield stacked outputs.

Claude

An LSTM cell allows the network to pick up such pertinent information and save it, injecting it back into the model when necessary.

While a basic RNN has only a simple activation function, an LSTM cell has four. Three sigmoid activation functions output numbers between 0 and 1, and lead to a pointwise multiplication gate. This gate determines whether information should enter the cell or not—0 means no information enters, and 1 means all the information enters. These three cells are used to save pertinent information for later stages of the learning process. 

In the training process, the network learns what are the optimal values for the gates, or how much of the information should be retained to help the network make the most accurate prediction.

LSTM Cell

Dive into Deep Learning

The forget gate is
$$\Gamma_f=\sigma(W_f[a^{<t-1>}, x^{<t>}]+b_f)$$
, where $\sigma$ denotes the sigmoid activation function, $W_f$ is weight, $b_f$ is a bias term, $a$ denotes the hidden state, $t$ is the t-th time/neuron, and $[a^{<t-1>}, x^{<t>}]$ means $a^{<t-1>}$ and $x^{<t>}$ are concatenated together.
Then compute the update gate in two steps. First, the update gate is
$$\Gamma_u=\sigma(W_u[a^{<t-1>}, x^{<t>}]+b_u)$$
Second, the intermediate cell state candidate is
$$\tilde c^{<t>}=tanh(W_c[a^{<t-1>}, x^{<t>}]+b_c)$$
, where $tanh$ denotes the tanh activation function.
Using the results from formulas above, we can calculate the current cell state,
$$c^{<t>}=\Gamma_u*\tilde c^{<t>}+\Gamma_f*c^{<t-1>}$$
At last, the third gate, output gate is
$$\Gamma_o=\sigma(W_o[a^{<t-1>}, x^{<t>}]+b_o)$$
, and using the output gate and the current cell state, we can compute the current hidden state
$$a^{<t>}=\Gamma_o*tanh(c^{<t>})$$

Math behind LSTMs

In a simple word,  word embeddings are numerical representations of texts.
In detail, when computer treat differents words, they will turn them into some atomic symbols. e.g. cat might be a, dogs might be b. But how can computer find the relationship between them? That's what word embeddings are. They willl build some relationships between words.

Word Embeddings

https://www.bioinf.jku.at/publications/older/2604.pdf

Long Short-Term Memory (LSTM) Paper

Gates are made use of by LSTM units to control the flow of information into and out of the units that comprise the network layers.
- The Forget Gate deletes information from the context that is no longer needed
- The Add Gate selects the information to add to the current context
- The Output Gate decides what information is required for the current hidden state

Gates in LSTMs

Given the appropriate weights for the various gates, an LSTM cell accepts three inputs at each time step: the memory cell internal state (context vector) and the hidden state from the previous time step, along with the current input vector. The cell then produces updated memory cell internal state and hidden state vectors as output. Crucially, only the hidden state is passed into the output layer or forwarded to subsequent layers in a stacked RNN, while the memory cell internal state remains entirely internal to the LSTM unit and is never exposed outside it.

Inputs and Outputs of LSTM Cells

The neural units used in LSTMs are obviously much more complex than those used in basic feedforward networks. Fortunately, this complexity is encapsulated within the basic processing units, allowing us to maintain modularity and easily experiment with different architectures. The figure below demonstrates this by illustrating the inputs and outputs associated with each kind of unit.

The increased complexity of the LSTM units is encapsulated within the unit itself. The only additional external complexity for the LSTM over the basic recurrent unit is the presence of the additional context vector as an input and output.

The Encapsulated Complexity of LSTMs

JAX lax.scan in RNN Iteration

When implementing a Long Short-Term Memory (LSTM) network, processing a new sequence without a prior state requires the explicit initialization of two distinct state variables. Both the hidden state ($$H$$) and the memory cell internal state ($$C$$) must be initialized, typically as zero tensors with a shape dictated by the batch size and the number of hidden units. This dual-state initialization provides a neutral starting point for the gating mechanisms and the input node before the first sequence token is processed.

Learn Before

Related