DOCSEARCHbyAlgolia
Submit Docs

PyxLSTM: Extended Long Short-Term Memory Implementation

repository·main·Indexed 43 days ago

https://github.com/muditbhargava66/pyxlstm

An efficient and extensible Python implementation of the xLSTM architecture (version 2.0.0) for sequence modeling tasks such as language modeling and text generation. It includes implementations of sLSTM (scalar LSTM) with exponential gating and mLSTM (matrix LSTM) for parallel computation on GPUs, along with utilities for tokenization, dataset preparation, and training loops.

Tokens
2.8K
Snippets
10
Records
13
Agent score
55%

What's inside PyxLSTM

  1. What is mLSTM (Matrix Long Short-Term Memory)?

    main

    mLSTM is a variant of the LSTM architecture that replaces the traditional scalar memory cell state with a matrix memory structure ($C$). It uses a covariance update rule and represents input information as key-value pairs ($k, v$) to enhance storage capacity.

    Key components include:

    • Input Gate ($i$): Controls input flow into the memory cell.
    • Forget Gate ($f$): Determines how much previous memory to retain.
    • Output Gate ($o$): Controls information flow from the cell to the hidden state.
    • Cell State ($C$): A matrix storing long-term memory.
    • Key ($k$) and Value ($v$): Represent input information for memory updates.
    • Query ($q$): Used to retrieve information from the memory matrix.

    Because it relies on matrix operations rather than sequential scalar updates, mLSTM is highly suitable for parallel computation on hardware accelerators like GPUs.

  2. How sLSTM architecture and gating work

    main

    The sLSTM architecture improves upon traditional LSTMs by introducing exponential gating for the input and forget gates, allowing for more fine-grained control over memory updates.

    Key components include:

    • Input Gate ($i$): Controls information flow into the memory cell using exponential gating.
    • Forget Gate ($f$): Determines how much previous memory to retain using exponential gating.
    • Output Gate ($o$): Controls information flow from the cell to the hidden state.
    • Cell State ($c$): Stores long-term memory.
    • Hidden State ($h$): The output at each time step.
    • Normalizer State ($n$): Used to scale the cell state updates.

    This architecture uses memory mixing, where the previous hidden state modulates the gates and cell state updates.

  3. PyxLSTM core modules and structure

    main

    The library is organized into several key modules:

    • xLSTM package:
      • slstm.py: Implementation of the sLSTM (scalar LSTM) module.
      • mlstm.py: Implementation of the mLSTM (matrix LSTM) module.
      • block.py: Implementation of xLSTM blocks (supports pre and post up-projection).
      • model.py: High-level xLSTM model definition.
    • utils package:
      • config.py: Configuration management.
      • logging.py: Logging setup.
      • utils.py: Miscellaneous utility functions.

  4. Prepare data for xLSTM training

    main

    Before training, text data must be tokenized and converted into sequences. Use the Tokenizer class to handle encoding/decoding and the LanguageModelingDataset class to format the data for training and validation.

    Key components:

    • Tokenizer: Requires a vocab_file to initialize.
    • LanguageModelingDataset: Takes the raw data, the initialized tokenizer, and a max_length parameter.
    from xLSTM.data import LanguageModelingDataset, Tokenizer

# Initialize tokenizer
tokenizer = Tokenizer(vocab_file)

# Create dataset
train_dataset = LanguageModelingDataset(train_data, tokenizer, max_length)
valid_dataset = LanguageModelingDataset(valid_data, tokenizer, max_length)

  5. Install PyxLSTM for development

    main

    To install the package in editable mode from the source repository, follow these steps:

    1. Clone the repository.
    2. Install requirements.
    3. Install the package in editable mode (-e).
    git clone https://github.com/muditbhargava66/PyxLSTM.git
cd PyxLSTM
pip install -r requirements.txt
pip install -e .

  6. Run unit tests in PyxLSTM

    main

    To verify the installation and ensure all modules (sLSTM, mLSTM, blocks, and model) are working correctly, run the unit tests using the unittest module from the root of the cloned repository.

    python -m unittest discover tests

  7. Implement a basic xLSTM training loop

    main

    The training process follows a standard PyTorch-style loop. You should move the model to the appropriate device using get_device() from xLSTM.utils and use CrossEntropyLoss for the criterion, ensuring you set ignore_index to tokenizer.pad_token_id to avoid calculating loss on padding tokens.

    Steps:

    1. Move model to device.
    2. Initialize optimizer (e.g., Adam) and criterion.
    3. Iterate through epochs and batches.
    4. Perform forward pass, calculate loss, backpropagate, and update weights.
    import torch
from xLSTM.utils import get_device

device = get_device()
model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
criterion = torch.nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)

for epoch in range(num_epochs):
    for batch in train_dataloader:
        inputs, targets = batch
        inputs, targets = inputs.to(device), targets.to(device)
        
        optimizer.zero_grad()
        outputs, _ = model(inputs)
        loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))
        loss.backward()
        optimizer.step()

  8. Install PyxLSTM

    main

    You can install the library using pip for standard usage or include development dependencies for testing and contributing.

    # Standard installation
pip install PyxLSTM

# Installation with development dependencies
pip install PyxLSTM[dev]

  9. Use PyxLSTM for language modeling

    main

    This example demonstrates the standard workflow for training an xLSTM model: loading configuration, initializing the tokenizer and dataset, creating the model, and running the training loop.

    import torch
from xLSTM.model import xLSTM
from xLSTM.data import LanguageModelingDataset, Tokenizer
from xLSTM.utils import load_config, set_seed, get_device
from xLSTM.training import train

# Load configuration
config = load_config("path/to/config.yaml")
set_seed(config.seed)
device = get_device()

# Initialize tokenizer and dataset
tokenizer = Tokenizer(config.vocab_file)
train_dataset = LanguageModelingDataset(config.train_data, tokenizer, config.max_length)

# Create xLSTM model
model = xLSTM(len(tokenizer), config.embedding_size, config.hidden_size,
              config.num_layers, config.num_blocks, config.dropout,
              config.bidirectional, config.lstm_type)
model.to(device)

# Train the model
optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
criterion = torch.nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
train(model, train_dataset, optimizer, criterion, config, device)

  10. Evaluate and save xLSTM model checkpoints

    main

    To monitor performance and prevent data loss, use the utility functions provided in xLSTM.utils for evaluation and checkpointing.

    • evaluate(model, dataloader, criterion, device): Calculates the loss on a validation dataset.
    • save_checkpoint(model, optimizer, checkpoint_path): Saves the current state of the model and optimizer to a specified path.
    from xLSTM.utils import evaluate, save_checkpoint

# Evaluate on validation set
valid_loss = evaluate(model, valid_dataloader, criterion, device)

# Save checkpoint
save_checkpoint(model, optimizer, checkpoint_path)

  11. Use the sLSTM class for scalar long short-term memory

    main

    The sLSTM (Scalar Long Short-Term Memory) is a variant of the LSTM architecture that uses exponential gating and memory mixing to improve state tracking. You can instantiate it by specifying the input size, hidden size, number of layers, and dropout rate. The model performs a forward pass that returns both the outputs and the hidden states.

    from xLSTM import sLSTM

# Create an sLSTM instance
slstm = sLSTM(input_size, hidden_size, num_layers, dropout)

# Forward pass
outputs, hidden_states = slstm(inputs)

  12. Use the mLSTM class for sequence modeling

    main

    You can use the mLSTM class from the xLSTM package to build models that utilize matrix-based memory. The class requires parameters for input size, hidden size, number of layers, and dropout. Once initialized, you can perform a forward pass by passing input tensors to the instance.

    from xLSTM import mLSTM

# Create an mLSTM instance
# Parameters: input_size, hidden_size, num_layers, dropout
mlstm = mLSTM(input_size, hidden_size, num_layers, dropout)

# Forward pass
# Returns: outputs, hidden_states
outputs, hidden_states = mlstm(inputs)