DOCSEARCHbyAlgolia
Submit Docs

minGRU-pytorch

repository·main·Indexed 43 days ago

https://github.com/lucidrains/mingru-pytorch

A log-space numerically stable implementation of the minGRU architecture in PyTorch (v0.2.1). It supports parallel sequence processing via associative scans and sequential token-by-token processing for autoregressive generation. The package includes the minGRU and minLSTM classes, a hybrid architecture combining minGRU with multi-head attention (minGRUAttnHybrid), and language model implementations including hybridLM and minLM.

Tokens
4.5K
Snippets
15
Records
18
Agent score
54%

What's inside minGRU-pytorch

  2. How minGRU handles parallel vs sequential processing

    main

    The minGRU module automatically switches between two processing modes based on the input sequence length:

    1. Sequential Mode (seq_len == 1): Used for step-by-step inference. It uses a linear interpolation (torch.lerp) between the prev_hidden and the current hidden state, gated by a sigmoid.
    2. Parallel Mode (seq_len > 1): Used for efficient training on sequences. It utilizes a log-space associative scan (heinsen_associative_scan_log) to compute the hidden states in parallel, which is more computationally efficient for long sequences than step-by-step recurrence.

    When prev_hidden is provided in parallel mode, the scan is prepended with the log of the previous hidden state to maintain continuity.

  3. Use minGRU in parallel mode

    main

    For efficient processing of entire sequences at once, pass the full sequence tensor to the minGRU module. This is the standard parallel implementation.

    Input shape: (batch, sequence_length, dim) Output shape: (batch, sequence_length, dim)

    import torch
from minGRU_pytorch import minGRU

min_gru = minGRU(dim = 512)

x = torch.randn(2, 1024, 512)

out = min_gru(x)

assert x.shape == out.shape

  4. Use minGRU in sequential mode

    main

    To process tokens one by one (e.g., for autoregressive generation), you can use the module sequentially by passing the prev_hidden state.

    When calling the module in sequential mode, use the argument return_next_prev_hidden = True to retrieve the updated hidden state for the next step. The sequential output should match the parallel output (within numerical tolerance).

    import torch
from minGRU_pytorch import minGRU

min_gru = minGRU(dim = 512, expansion_factor = 1.5)
x = torch.randn(1, 2048, 512)

# parallel mode for reference
parallel_out = min_gru(x)[:, -1:]

# sequential mode
prev_hidden = None
for token in x.unbind(dim = 1):
    sequential_out, prev_hidden = min_gru(token[:, None, :], prev_hidden, return_next_prev_hidden = True)

assert torch.allclose(parallel_out, sequential_out, atol = 1e-4)

  5. Initialize the minGRU class

    main

    The minGRU class implements a log-space numerically stable version of the minGRU architecture.

    Parameters:

    • dim: The dimensionality of the input/output.
    • expansion_factor: (Optional) The expansion factor used in the architecture.
    from minGRU_pytorch import minGRU

# Basic initialization
min_gru = minGRU(dim = 512)

# Initialization with expansion factor
min_gru = minGRU(dim = 512, expansion_factor = 1.5)

  6. Configure minGRUAttnHybrid parameters

    main

    When initializing minGRUAttnHybrid, you can configure the following parameters:

    • dim: The input and output dimension of the module.
    • dim_head: The dimension of each attention/RNN head. Defaults to 64.
    • heads: The number of heads to use. Defaults to 8.
    • learned_mix: A boolean indicating whether to use a learned mixing coefficient for the two branches. If True, a to_mix layer is added to compute per-token/per-head weights. Defaults to True.

  7. Initialize the minLM model

    main

    The minLM class implements a language model architecture based on minGRU or minLSTM. It supports causal depth-wise convolutions and configurable expansion factors.

    Note on Caching: If enable_conv is set to True, the model uses CausalDepthWiseConv1d, which prevents the use of prev_hiddens for recurrent decoding (caching is not supported in the convolutional version).

    Parameters:

    • num_tokens: Vocabulary size.
    • dim: Embedding and hidden dimension.
    • depth: Number of layers.
    • ff_mult: Multiplier for the FeedForward expansion (default: 4).
    • expansion: Expansion factor for the minGRU/minLSTM layer (default: 1.5).
    • conv_kernel_size: Kernel size for the causal depth-wise convolution (default: 3).
    • use_lstm: If True, uses minLSTM instead of minGRU (default: False).
    • enable_conv: If True, enables the CausalDepthWiseConv1d layer in each block (default: False).
    • dropout: Dropout probability (default: 0.0).
    from minGRU_pytorch.minLM import minLM

model = minLM(
    num_tokens = 5000, 
    dim = 256, 
    depth = 6, 
    enable_conv = True
)

  9. Initialize hybridLM with configuration parameters

    main

    The hybridLM constructor accepts the following keyword-only arguments:

    • num_tokens: The size of the vocabulary (integer).
    • dim: The embedding dimension (integer).
    • depth: The number of transformer-like layers (integer).
    • ff_mult: The multiplier for the FeedForward inner dimension (default: 4).
    • heads: The number of attention heads for the hybrid block (default: 8).
    • dim_head: The dimension of each attention head (default: 64).

  10. Use CausalDepthWiseConv1d for causal convolutions

    main

    The CausalDepthWiseConv1d module applies a depth-wise convolution followed by a 1x1 convolution. It uses padding to ensure the convolution is causal (the output at position $t$ only depends on positions $\le t$).

    Arguments:

    • dim: Input and output dimension.
    • kernel_size: Size of the depth-wise convolution kernel.
    from minGRU_pytorch.minLM import CausalDepthWiseConv1d
import torch

conv = CausalDepthWiseConv1d(dim=256, kernel_size=3)
x = torch.randn(8, 100, 256) # (batch, seq_len, dim)
output = conv(x)

  11. Perform forward pass with minLM

    main

    The forward method of minLM processes input token IDs and can return logits, loss, or hidden states for recurrent decoding.

    Arguments:

    • x: Input tensor of shape (batch, sequence_length).
    • return_loss: If True, the method expects x to contain both tokens and labels (where labels are shifted by one). It returns the cross-entropy loss.
    • return_prev_hiddens: If True, returns the hidden states from the minGRU/minLSTM layers, which can be used for subsequent steps in recurrent decoding.
    • prev_hiddens: A list of previous hidden states for each layer, used when performing step-by-step recurrent decoding.

    Returns:

    • If return_loss is False and return_prev_hiddens is False: Returns logits of shape (batch, sequence_length, num_tokens).
    • If return_loss is True: Returns loss (scalar).
    • If return_prev_hiddens is True: Returns logits or (logits, next_prev_hiddens) depending on return_loss.
    # Standard forward pass
logits = model(x)

# Forward pass for training with loss
loss = model(x, return_loss=True)

# Forward pass for recurrent decoding (step-by-step)
# Note: enable_conv must be False to use prev_hiddens
logits, next_hiddens = model(x, return_prev_hiddens=True, prev_hiddens=my_prev_hiddens)

  12. Initialize and use the minGRU module

    main

    The minGRU class implements the minGRU architecture as described in arXiv:2410.01201v1. It can operate in both parallel mode (for training/long sequences) and sequential mode (for single-step inference).

    Parameters

    • dim: The input dimension.
    • expansion_factor: A multiplier for the internal dimension. Defaults to 1.0.
    • proj_out: If provided, determines the output projection. If None, it defaults to expansion_factor != 1. (meaning if you expand the dimension, it will project back to the original dim).

    Forward Pass

    forward(x, prev_hidden = None, return_next_prev_hidden = False)

    • x: Input tensor of shape (batch, seq_len, dim).
    • prev_hidden: The previous hidden state of shape (batch, 1, dim). If provided, the model performs stateful updates.
    • return_next_prev_hidden: If True, returns a tuple containing the output and the updated hidden state for the next step.
    from minGRU_pytorch import minGRU
import torch

# Initialize
model = minGRU(dim = 128, expansion_factor = 2.0)

# Input: (batch, seq_len, dim)
x = torch.randn(32, 10, 128)

# Forward pass
out = model(x)

# Forward pass with state tracking
out, next_hidden = model(x, return_next_prev_hidden = True)