Research model for Sparse Query Attention experiments - extension to Grouped Query Attention, that's also reducing the number of used query heads, instead of further reducing key/value heads count (up to Multi Query Attention). That approach results in huge computational complexity reduction and much faster training, while the performance stays on GQA level (almost unnoticeable decrease, when compared to GQA, and noticeable better than MQA).
This version is using an extreme 4x query heads reduction factor with a single key/value group to combine Multi Query Attention with Sparse Query Attention. However, further key/value head reduction doesn't bring any training time reduction, compared to xSQA and noticeably decreasing performance. So this solution is not recommended.
Research paper - arxiv.org/abs/2510.01817
This model was only trained for research purposes, on a small number of training steps.
Validation mean loss/accuracy:
Training time / time per batch:
SQA has reduced dimensions of query heads linear projection and output projection, which results in a little smaller model sizes:
Model requires our RxLM framework for training/inference. It's integrated with HuggingFace Hub and libraries. Components connected to SQA and classic transformers are free even for commercial usage, while Reactive Transformer components are free only for non-commercial usage (Reactive AI Framework License v1.0)
pip install rxnn torch transformers tokenizersimport torch
from rxlm.experimental.models import ExperimentalAttentionTransformer
from rxlm.transformers.sampler import Sampler, SampleDecoder
from rxlm.training.tokenizer import load_tokenizer_from_hf_hub
model = ExperimentalAttentionTransformer.from_pretrained('ReactiveAI/xSMQAT-m')
tokenizer = load_tokenizer_from_hf_hub('ReactiveAI/xSMQAT-m')
sampler = Sampler(model, torch.device('cuda' if torch.cuda.is_available() else 'cpu'), end_token_id=3)
sample = SampleDecoder(sampler, tokenizer)
# 0.1 and 0.9 are default values for temperature and top_p
generated = sample('Example model input for text generation...', temperature=0.1, top_p=0.9, max_seq_len=1024)
sample('Example model input for text generation - print streamed response...', temperature=0.1, top_p=0.9, max_seq_len=1024, print_stream=True)
pip install rxnn torch transformers tokenizers tensorboard (tensorboard is optional)import torch
from rxlm.experimental.models import ExperimentalAttentionTransformer
from rxlm.training.tokenizer import load_tokenizer_from_hf_hub
from rxlm.llm_training.dataset import AutoregressiveLMDataset
from rxlm.llm_training.supervised import AutoregressiveTrainer
from rxlm.training.callbacks import PrintLossCallback, PrintAccuracyCallback, TokenCounterCallback, ModelSaveCallback
from rxlm.training.scheduler import get_transformer_lr_scheduler
model = ExperimentalAttentionTransformer.from_pretrained('ReactiveAI/xSMQAT-m')
tokenizer = load_tokenizer_from_hf_hub('ReactiveAI/xSMQAT-m')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
batch_size = 128 # Require ~40GB GPU Memory (trained on L40S)
epochs = 1
gradient_acc_steps = 1
seq_len = 1024
vocab_size = 10_000
peak_lr = 5e-4 * gradient_acc_steps
train_dataset = AutoregressiveLMDataset.from_hf_hub('hf-dataset-id', 'subset', tokenizer=tokenizer, max_seq_len=seq_len) # split is 'train' by default
valid_dataset = AutoregressiveLMDataset.from_hf_hub('hf-dataset-id', split='validation', tokenizer=tokenizer, max_seq_len=seq_len)
dataset_len = len(train_dataset)
steps_per_epoch = int(dataset_len / batch_size - 1)
total_steps = int((epochs * steps_per_epoch) / gradient_acc_steps)
warmup_steps = int(0.25 * steps_per_epoch)
logs_dir = './tensorboard_logs' # require tensorboard `pip install tensorboard`
print_cb = PrintLossCallback(batches_per_epoch=steps_per_epoch)
count_cb = TokenCounterCallback()
acc_cb = PrintAccuracyCallback()
save_cb = ModelSaveCallback('./path/to/save', push_to_hub=True,
hub_model_id='your-model-id', private_repo=True,
push_checkpoint_weights=True, final_commit_message='Final commit message', hf_token=YOUR_HF_TOKEN)
trainer = AutoregressiveTrainer(model, device, dataset=train_dataset, validation_dataset=valid_dataset,
vocab_size=vocab_size, callbacks=[print_cb, acc_cb, count_cb, save_cb], use_amp=True,
dtype=torch.bfloat16, log_dir=logs_dir, gradient_accumulation_steps=gradient_acc_steps)
optimizer = torch.optim.AdamW(model.parameters(), lr=peak_lr, weight_decay=0.01)
scheduler = get_transformer_lr_scheduler(
optimizer,
warmup_steps=warmup_steps,
num_training_steps=total_steps
)
trainer(epochs=epochs, batch_size=batch_size, optimizer=optimizer, scheduler=scheduler)
According to experiment results, this variant has the worst performance and no gains in training time. It seems, that after some point, further key/value head reduction isn't a good idea and this variant is not recommended. For the fastest training time, xSQA has the same results and is more efficient.