Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Fully featured implementation of Routing Transformer
| Date | Stars |
|---|---|
| 2026-07-24 | 300 |
| 2026-07-25 | 300 |
| 2026-07-28 | 300 |
| 2026-07-30 | 300 |
| 2026-08-06 | 300 |
Today
— stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
## Routing Transformer
<img src="./routing_attention.png" width="500px"></img>
[](https://badge.fury.io/py/routing-transformer)
A fully featured implementation of <a href="https://arxiv.org/pdf/2003.05997.pdf">Routing Transformer</a>. The paper proposes using k-means to route similar queries / keys into the same cluster for attention.
[](https://colab.research.google.com/drive/1sw1Hjx3EllhKZh4nhJ3TIZ978HjKVUFQ?usp=sharing) 131k tokens
### Install
```bash
$ pip install routing_transformer
```
### Usage
A simple language model
```python
import torch
from routing_transformer import RoutingTransformerLM
model = RoutingTransformerLM(
num_tokens = 20000,
dim = 512,
heads = 8,
depth = 12,
max_seq_len = 8192,
causal = True, # auto-regressive or not
emb_dim = 128, # embedding factorization, from Albert
weight_tie = False, # weight tie layers, from Albert
tie_embedding = False, # multiply final embeddings with token weights for logits
dim_head = 64, # be able to fix the dimension of each head, making it independent of the embedding dimension and the number of heads
attn_dropout = 0.1, # dropout after attention
attn_layer_dropout = 0., # dropout after self attention layer
ff_dropout = 0.1, # feedforward dropout
layer_dropout = 0., # layer dropout
window_size = 128, # target window size of each cluster
n_local_attn_heads = 4, # number of local attention heads
reversible = True, # reversible networks for memory savings, from Reformer paper
ff_chunks = 10, # feed forward chunking, from Reformer paper
ff_glu = True, # use GLU variant in feedforward
pkm_layers = (4, 7), # specify layers to use product key memory. paper shows 1 or 2 modules near the middle of the transformer is best
pkm_num_keys = 128, # defaults to 128, but can be increased to 256 or 512 as memory allows
moe_layers = (3, 6), # specify which layers to use mixture of experts
moe_num_experts = 4, # number of experts in the mixture of experts layer, defaults to 4. increase for adding more parameters to model
moe_loss_coef = 1e-2, # the weight for the auxiliary loss in mixture of experts to keep expert usage balanced
num_mem_kv = 8, # number of memory key/values to append to each cluster of each head, from the 'All-Attention' paper. defaults to 1 in the causal case for unshared QK to work
use_scale_norm = False, # use scale norm, simplified normalization from 'Transformers without Tears' paper
use_rezero = False, # use Rezero with no normalization
shift_tokens = True # shift tokens by one along sequence dimension, for a slight improvement in convergence
).cuda()
x = torch.randint(0, 20000, (1, 8192)).long().cuda()
input_mask = torch.ones_like(x).bool().cuda()
y, aux_loss = model(x, input_mask = input_mask) # (1, 8192, 20000)
aux_loss.backward() # add auxiliary loss to main loss before backprop
```
A simple transformer
```python
import torch
from routing_transformer import RoutingTransformer
model = RoutingTransformer(
dim = 512,
heads = 8,
depth = 12,
max_seq_len = 8192,
window_size = 128,
n_local_attn_heads = 4
).cuda()
x = torch.randn(1, 8192, 512).cuda()
input_mask = torch.ones(1, 8192).bool().cuda()
y, aux_loss = model(x, input_mask = input_mask) # (1, 8192, 512)
aux_loss.backward() # add auxiliary loss to main loss before backprop
```
## Encoder Decoder
To use a full encoder, decoder, simply import the `RoutingTransformerEncDec` class. Save for the `dim` keyword, all other keywords will be either prepended with `enc_` or `dec_` for the encoder and decoder `RoutingTransformerLM` class respectively.
```python
import torch
from routing_transformer import RoutingTranExcerpt of 11,375 characters
Read on GitHubWould you bet a product on this? Bounded 0–100 and slow moving.
matched fp:48e3cffcda617af8, topic:deep-learning, topic:pytorch
matched fp:48e3cffcda617af8, topic:transformer, readme:mixture of experts