GRU

GRU

The Gated Recurrent Unit is an RNN that combines forget and input gates of the LSTM into a single update gate, and merges cell state and hidden state.

$$\begin{array}{rlrl}{z}_t& =\sigma (W_z{x}_t+U_z{h}_{t-1}+b_z)& & \text{(update gate)}\\ r_t& =\sigma (W_r{x}_t+U_r{h}_{t-1}+b_r)& & \text{(reset gate)}\\ {\tilde{h}}_t& =\tanh (W_h{x}_t+U_h(r_t\odot h_{t-1})+b_h)& & \text{(candidate state)}\\ h_t& =(1-z_t)\odot h_{t-1}+z_t\odot {\tilde{h}}_t& & \text{(new hidden state)}\end{array}$$

$x_t$: input vector at time (t)
$h_{t-1},h_t$: previous and current hidden states
$z_t,r_t$ … update and reset gates (values between 0 and 1)
$\sigma ,\tanh$: sigmoid, tanh
$\odot$ … elementwise multiplication
$W_{\ast }$ … input weight matrices
$U_{\ast }$ … hidden state weight (update) matrices
$b_{\ast }$ … bias vectors

The notation here summarizes $W_{\ast }$ and $U_{\ast }$ into a single weight matrix each, and concatenates the input and hidden state vectors. Biases are like so often omitted for simplicity.

This matches how it’s actually implemented, doing 2 instead of 6 matmuls:

class GRUCellDiagram(nn.Module):
	def __init__(self, input_size, hidden_size):
		super().__init__()
		self.gates = nn.Linear(input_size + hidden_size, 2 * hidden_size) # W_z and W_r
		self.cand  = nn.Linear(input_size + hidden_size, hidden_size) # W 
 
	def forward(self, x_t, h_prev): # (batch, input_size), (batch, hidden_size)
		hx = torch.cat([h_prev, x_t], dim=-1)
		r_t, z_t = self.gates(hx).chunk(2, dim=-1)
		r_t = torch.sigmoid(r_t)
		z_t = torch.sigmoid(z_t)
 
		rhx   = torch.cat([r_t * h_prev, x_t], dim=-1)
		h_hat = torch.tanh(self.cand(rhx))
 
		h_t = (1.0 - z_t) * h_prev + z_t * h_hat
		return h_t

https://colah.github.io/posts/2015-08-Understanding-LSTMs/