WERSA
2000년 도입
# WERSA: Wavelet-Enhanced Random Spectral Attention
This repository provides the official implementation of WERSA, a novel attention mechanism with linear O(n) time complexity, designed to scale Transformer models to very long sequences without a performance trade-off.
Our paper, "Scaling Attention to Very Long Sequences in Linear Time with Wavelet-Enhanced Random Spectral Attention (WERSA)", is available on arXiv:2507.08637.
## 🔬 The Science Behind WERSA
Standard attention mechanisms have a quadratic (O(n²)) complexity that makes processing long sequences impractical. WERSA solves this by combining several powerful principles to achieve linear (O(n)) efficiency while maintaining high performance.
- Multi-Resolution Analysis: Uses Haar wavelet transforms to decompose the input into multiple scales, capturing both local details and global context.
- Adaptive Filtering: An MLP generates input-dependent filters and learnable scale_weights modulate each wavelet level, allowing the model to dynamically prioritize the most informative frequency components.
- Linear Complexity via Random Features: Uses random feature projection to approximate the softmax kernel, avoiding the computation of the full quadratic attention matrix.
## ⚙️ Installation
First, ensure you have PyTorch and Hugging Face Transformers installed. Then, install the wersa package directly from this repository.
``bash
# 1. Install core dependencies (example for CUDA 12.1)
pip install torch --index-url https://download.pytorch.org/whl/cu121
pip install transformers
# 2. Install the WERSA package from this repository
pip install git+https://github.com/vincenzodentamaro/wersa.git
`
## 🚀 Quickstart: Building a Qwen-like Model with WERSA
You can easily build a Qwen-style causal language model with WERSA attention by importing the WersaConfig and WersaForCausalLM classes from the package.
### Building an 8B Parameter Model
This snippet creates an ~8B parameter model with a configuration similar to state-of-the-art models like Qwen2-7B.
`python
from wersa import WersaConfig, WersaForCausalLM
from transformers import AutoTokenizer
# Load a compatible tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B")
# Define the configuration for the 8B model
config_8b = WersaConfig(
vocab_size=len(tokenizer),
pad_token_id=tokenizer.pad_token_id,
hidden_size=4096,
num_hidden_layers=32,
num_attention_heads=32,
intermediate_size=11008,
max_position_embeddings=4096
)
# Instantiate the model
model_8b = WersaForCausalLM(config_8b)
print(f"8B Model created with ~{model_8b.num_parameters() / 1e9:.2f}B parameters.")
`
### Building a 0.6B Parameter Model
This snippet creates a smaller ~0.6B parameter model, perfect for faster experiments or deployment on more constrained hardware.
`python
from wersa import WersaConfig, WersaForCausalLM
from transformers import AutoTokenizer
# Load a compatible tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B")
# Define the configuration for the 0.6B model
config_0_6b = WersaConfig(
vocab_size=len(tokenizer),
pad_token_id=tokenizer.pad_token_id,
hidden_size=1024,
num_hidden_layers=24,
num_attention_heads=16,
intermediate_size=2816,
max_position_embeddings=1024
)
# Instantiate the model
model_0_6b = WersaForCausalLM(config_0_6b)
print(f"0.6B Model created with ~{model_0_6b.num_parameters() / 1e9:.2f}B parameters.")
`
## 📖 Training and Examples
This repository includes complete scripts to demonstrate how to pre-train these models from scratch and test their generation capabilities.
- train_and_generate_1b.py: A full example for training a ~1B parameter model.
- train_and_generate_8b.py: A full example for training the 8B parameter model.
## 📜 Citation
If you find WERSA useful in your research, please consider citing our paper:
`bibtex
@misc{dentamaro2025scaling,
title={Scaling Attention to Very Long Sequences in Linear Time with Wavelet-Enhanced Random Spectral Attention (WERSA)},
author={Vincenzo Dentamaro},
year={2025},
eprint={2507.08637},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
``
## 📄 License
This project is licensed under the Apache License 2.0.