kl-div-attention 0.1.2

Attention with QK distance using KL divergence

kl-div-attention

Just another attention variant, where the query key distance uses KL divergence, for testing Gitlab workflow

Install

pip install kl-div-attention

Usage

import torch
from kl_div_attention import KLDivAttention

attn = KLDivAttention(
    dim = 512,
    heads = 8,
    dim_head = 64,
    causal = True,
    prenorm = True, 
    fused_mode = 'flash'    # use fused triton flash attention
).cuda()

tokens = torch.randn(1, 1024, 512).cuda()

out = attn(tokens) + tokens

assert out.shape == tokens.shape

out.sum().backward()

Training

You can train a small transformer on Enwik8 with KL-divergence attention using the provided training script.

Eg. to use the fused Triton kernel

uv run train_enwik8.py --fused-mode flash