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

pip3 install -U jax-transformer

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.

Usage

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

import jax
from jax_transformer.main import transformer_decoder, causal_mask

# Example usage
batch_size = 2
seq_len = 10
dim = 64
heads = 8
d_ff = 256
depth = 6

# Random input tokens
x = jax.random.normal(
    jax.random.PRNGKey(0), (batch_size, seq_len, dim)
)
rng = jax.random.PRNGKey(42)
# Generate causal mask
mask = causal_mask(seq_len)

# Run through transformer decoder
out = transformer_decoder(
    x=x,
    mask=mask,
    depth=depth,
    heads=heads,
    dim=dim,
    d_ff=d_ff,
    dropout_rate=0.1,
    rng=rng,
)


print(out.shape)  # Should be (batch_size, seq_len, dim)

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.

Citation

@article{JaxTransformer,
    author={Kye Gomez},
    title={Jax Transformer},
    year={2024},
}