This repo hosts the framework for development, pre-training and evaluation of efficient neural network architectures for language modeling, and sequence processing in general.

What means "efficient" sequence-mixing architecture?

• Capable of running in sub-quadratic, preferably linear time O(N) w.r.t sequence length N.
• Having a competitive processing speed not only for extremely long sequences but for short to moderate context sizes (e.g., it should be comparable in throughput to softmax attention/ Transformer at N=128).

Currently, the framework allows for off-the-shelf reuse or rapid prototyping of the following architectures and their compositions:

• DANet and DenseAttention (see below)
• Transformer (softmax attention)
• Numerous Linear Attention architectures which conform to framework described in Transformers are RNNs
• Various locality patterns: Local Attention (also known as Block-Sparse Attention), Shifted-Local Attention, Sliding Window Attention (SWA), which are applicable not only for softmax attention, but for other sequence mixers.
• Architectures implemented in HuggingFace's Transformers library.

Coming soon:

• Support for architectures implemented in FLA framework
• S4 and other linear-time-invariant SSMs.

Tasks supported by the framework:

• Masked Language (or Sequence) Modeling
• Causal Language (Sequence) Modeling
• Classification (standalone or paired with sequence modeling)
• Regression

Notable features:

• Use mixtures of any datasets, either from HuggingFace or stored locally, in streaming mode;
• Use any patterns of local and global layers for most models
• Create any types of hybrid architectures (coming soon)
• Analyze rich model metrics during pre-training (e.g., distributions/ norms of activations, weight and gradients) and measure speed of pre-training and evals.
• Use distributed mode or mixed-precision training seamlessly via DeepSpeed

The framework is currently under active development and will be used in future research.

This repository also hosts the code of the official implementation and the experiments for the paper "MatMuls are Enough for Efficient and Performant Linear-Time Attention".

ICML 2025 Workshops version of the paper:

• LCFM (short format) https://openreview.net/forum?id=yLh58rr4JX
• ES-FoMo III (extended version) https://openreview.net/forum?id=RttNumxC1t

(Also available locally)

DANet-BERT models on HuggingFace: to be added soon

In this research, we propose DenseAttention and DenseAttention Network (DANet) as a replacement for softmax self-attention and Transformer block, respectively. Key advantages of the new architecture include:

• $O(N)$ time and space complexity with respect to sequence length $N$;
• Speed vs Transformers: DANet is faster than low-level FlashAttention implementation even on small sequences despite being written in plain PyTorch;
• Faster speed over other linear-time sequence processing algorithms;
• Compatibility and accessibility: DenseAttention does not require low-level CUDA code, runs on every device where PyTorch can be installed, and works well with both bf16 and older fp16 half-precision formats.

DenseAttention achieves better speed, linear complexity, and computational efficiency without compromising the modeling performance. This is exemplified by, among other experiments, DANet-BERT-Large LM pre-training on approximately 500B tokens and fine-tuning on the GLUE benchmarks, which yields best results out of all models of comparable size trained on sub-trillion token count data. It's also supported by validation on the LRA suite of benchmarks where it outperforms all previous Transformer-based architectures at least by 5%.

DenseAttention is a novel self-attention mechanism which eliminates softmax and does not introduce any replacements. It also merges $W_Q$ and $W_K$ projection matrices into a single parameter. Remarkably, it is composed entirely of dense MatMuls.

The whole DANet further simplifies Transformer module by reducing the number of other element-wise operations and merging projection matrices. To ensure numerical stability, We replace standard LayerNorm with MaxNormActivation, and for adding sharp focus on nearby tokens in extremely long contexts, we introduce local and shifted local DenseAttention layers designed to complement standard layers with global receptive field. More details are provided in the paper.

This is an ongoing project. Some parts of the code and configs are subject to change. For the most current version and new features, please checkout dev branch.

The code for DenseAttention Network, its constituents and models built upon it can be found in src/ directory.

All settings for DenseAttention and DANet are documented in HuggingFace-like model config file.

DANet can serve as a drop-in replacement for Transformer modules. To start building your own models using DANet, copy src directory into your project and import required layers:

import torch
from src.danet_layers import DANetLayer, DANetLayerWithLocalAttention
from src.model_config import ModelConfig

SIZE=128
config = ModelConfig(hidden_size=SIZE)
danet_layer = DANetLayer(config) # or DANetLayerWithLocalAttention(config)
inputs = torch.randn(8, 1024, SIZE) # batch size, sequence length, model dimension
outputs = danet_layer(inputs)

For examples on how to incorporate various types of positional embeddings into DANet with and without local attention, see DANetEncoder in modeling.

For reproducing the experiments from the paper, it's recommended to create Docker container with preinstalled Cuda or a dedicated conda environment. To start, run:

git clone https://github.com/andrewargatkiny/dense-attention.git
cd dense-attention && pip install -r requirements.txt
export BASE_DATA_DIR="path/to/desired/data/directory"

To prepare the data, use scripts in the data_preparation/ directory. It can take a long time.

To reproduce an experiment, or simply train / evaluate a model, launch corresponding ds_train_*.sh script from the configs/ directory which contains exact configurations for the experiments in the paper. Under the hood, it launches deepspeed_train.py with a preconfigured set of arguments. Additionally, you can override any configuration parameter directly from the command line using the --override flag, without modifying the original config files. To construct your own pipeline or to learn about each argument's effect, please refer to train_arguments.py script and Using Configs to Run Experiments.

It’s recommended to use ClearML open-source ML experiments tracking system for the training. You can use it in the cloud up to a small storage limit or install on your server, for free. Here are the instructions on how to install it on your system and how to use it. ClearML is enabled by default, but it's also possible to use Weights & Biases or Tensorboard by providing --wandb or --tensorboard argument in deepspeed_train.py params of the training script.

If you use DenseAttention in research or production, or otherwise find it useful, please cite it as:

@inproceedings{
argatkiny2025matmuls,
title={MatMuls are Enough for Efficient and Performant Linear-Time Attention},
author={Andrew Argatkiny and Ilya Makarov},
booktitle={ES-FoMo III: 3rd Workshop on Efficient Systems for Foundation Models},
year={2025},
url={https://openreview.net/forum?id=RttNumxC1t}
}