jax-transformer 0.0.1

Jax Transformer - Jax

Jax Transformer

This repository demonstrates how to build a Decoder-Only Transformer with Multi-Query Attention in JAX. Multi-Query Attention is an efficient variant of the traditional multi-head attention, where all attention heads share the same key-value pairs, but maintain separate query projections.

Table of Contents

• Overview
• Key Concepts
• Installation
• Usage
• Code Walkthrough
  • Multi-Query Attention
  • Feed-Forward Layer
  • Transformer Decoder Block
  • Causal Masking
• Running the Transformer Decoder
• Contributing
• License

Overview

This project is a tutorial for building Transformer models from scratch in JAX, with a specific focus on implementing Decoder-Only Transformers using Multi-Query Attention. Transformers are state-of-the-art models used in various NLP tasks, including language modeling, text generation, and more. Multi-Query Attention (MQA) is an optimized version of multi-head attention, which reduces memory and computational complexity by sharing key and value matrices across all heads.

Key Concepts

• Multi-Query Attention: Shares a single key and value across all attention heads, reducing memory usage and computational overhead compared to traditional multi-head attention.
• Transformer Decoder Block: A core component of decoder models, which consists of multi-query attention, a feed-forward network, and residual connections.
• Causal Masking: Ensures that each position in the sequence can only attend to itself and previous positions to prevent future token leakage during training.

Installation

To get started with this project, clone the repository and install the dependencies:

git clone https://github.com/your-username/jax-transformer-multi-query-attention.git
cd jax-transformer-multi-query-attention

# Install required libraries
pip install -r requirements.txt

Requirements

• JAX: A library for high-performance machine learning research. Install JAX with GPU support (optional) by following the instructions on the JAX GitHub page.
• Numpy: For efficient numerical operations.
• Matplotlib (optional): If you wish to visualize attention weights or other model outputs.

Usage

After installing the dependencies, you can run the model on random input data to see how the transformer decoder works:

import jax
from transformer_decoder import transformer_decoder, causal_mask

# Define model dimensions
batch_size = 2
seq_len = 10
d_model = 64
num_heads = 8
d_ff = 256
num_layers = 6

# Random input tensor (representing tokens in the sequence)
x = jax.random.normal(jax.random.PRNGKey(0), (batch_size, seq_len, d_model))

# Generate causal mask for auto-regressive behavior
mask = causal_mask(seq_len)

# Run the input through the decoder
output = transformer_decoder(x, mask, num_layers=num_layers, num_heads=num_heads, d_model=d_model, d_ff=d_ff)

print(output.shape)  # Output: (batch_size, seq_len, d_model)

Code Walkthrough

This section explains the key components of the model in detail.

Multi-Query Attention

The Multi-Query Attention mechanism replaces the traditional multi-head attention by sharing the same set of key-value pairs for all heads while keeping separate query projections. This drastically reduces the memory footprint and computation.

def multi_query_attention(query, key, value, mask):
    ...

Feed-Forward Layer

After the attention mechanism, the transformer applies a two-layer feed-forward network with a ReLU activation in between. This allows the model to add depth and capture complex patterns.

def feed_forward(x, d_ff):
    ...

Transformer Decoder Block

The Transformer Decoder Block combines the multi-query attention mechanism with the feed-forward network and adds residual connections and layer normalization to stabilize the learning process. It processes sequences in a causal manner, meaning that tokens can only attend to previous tokens, which is crucial for auto-regressive models (e.g., language models).

def transformer_decoder_block(x, key, value, mask, num_heads, d_model, d_ff):
    ...

Causal Masking

The Causal Mask ensures that during training or inference, tokens in the sequence can only attend to themselves or previous tokens. This prevents "future leakage" and is crucial for tasks such as language modeling and text generation.

def causal_mask(seq_len):
    ...

Running the Transformer Decoder

To run the decoder model, execute the following script:

python run_transformer.py

The model takes random input and runs it through the Transformer decoder stack with multi-query attention. The output shape will be (batch_size, seq_len, d_model).

Contributing

Contributions are welcome! If you'd like to contribute, please fork the repository and submit a pull request with your improvements. You can also open an issue if you find a bug or want to request a new feature.

License

This project is licensed under the MIT License - see the LICENSE file for details.