QSoftmax

Quantum-Inspired Softmax Alternative for Attention Mechanisms

Abstract

QSoftmax is a novel attention normalization mechanism inspired by quantum mechanical principles. Instead of using exponential normalization (softmax), QSoftmax encodes vectors as quantum states and computes attention weights through interference patterns—where aligned phases produce constructive interference (high attention) and misaligned phases produce destructive interference (zero attention).

This approach provides natural sparsity without explicit top-k selection and offers a theoretically grounded alternative to address the attention dilution problem in long-context transformers.

Motivation

The Attention Dilution Problem

Standard softmax attention suffers from a fundamental scaling issue as context length increases:

                    exp(score_i)
attention_i = ─────────────────────
               Σⱼ exp(score_j)

From extreme value theory, the maximum score among L random values grows only as O(√log L), while the denominator grows as O(L). This causes attention to dilute—even highly relevant tokens receive vanishingly small attention weights at long context lengths.

The QSoftmax Approach

QSoftmax replaces exponential normalization with quantum interference:

attention_i = |⟨query|key_i⟩|²

Where query and key are encoded as quantum states with complex amplitudes. This provides:

Natural sparsity: Destructive interference produces exact zeros
Phase-based matching: Semantic similarity encoded in phase alignment
No explicit normalization denominator: Avoids the O(L) dilution term

Method

1. Quantum State Encoding

Each vector is encoded into a 4-dimensional quantum state (2-qubit analogy):

|ψ⟩ = α|00⟩ + β|01⟩ + γ|10⟩ + δ|11⟩

Where α, β, γ, δ are complex amplitudes:

Magnitude: Derived from vector energy (L2 norm of subgroups)
Phase: Derived from weighted positional information

Constraint: |α|² + |β|² + |γ|² + |δ|² = 1 (normalization)

2. Interference-Based Scoring

Attention scores are computed via the Born rule:

score(q, k) = |⟨query|key⟩|² = |Σᵢ conj(qᵢ) × kᵢ|²

Constructive interference: Phases align → amplitudes add → high score
Destructive interference: Phases oppose → amplitudes cancel → zero score

3. Sparse Normalization

# Apply threshold (creates natural sparsity)
scores[scores < max(scores) * threshold] = 0

# Amplitude normalization (sqrt, not exp)
weights = sqrt(scores)
weights = weights / sum(weights)

4. Optional Entanglement

Nearby keys can be "entangled" to capture local correlations:

entangle(key[i], key[i+1], strength=0.3)

Benchmark Results

Dashboard Overview

A) Needle-in-Haystack Performance

Measures ability to focus attention on a specific relevant token among distractors.

Method	Seq=64	Seq=128	Seq=256	Seq=512	Seq=1024
Sparsemax	1.000	1.000	1.000	1.000	1.000
Softmax	0.316	0.224	0.141	0.071	0.032
QSoftmax	3.5e-6	0.012	0.007	0.002	<1e-6

Note: Current QSoftmax implementation shows lower needle attention than expected. This indicates the interference threshold may be too aggressive, causing the needle to be filtered out. Parameter tuning is ongoing.

B) Attention Sparsity

Percentage of attention weights that are effectively zero.

Method	Seq=64	Seq=128	Seq=256	Seq=512	Seq=1024
Sparsemax	98.4%	99.2%	99.6%	99.8%	99.9%
QSoftmax	0%	1.6%	7.0%	15.2%	25.1%
Softmax	0%	0%	0%	0%	0%

QSoftmax achieves natural sparsity through destructive interference, scaling with sequence length.

C) Attention Entropy

Lower entropy indicates more focused attention distribution.

Method	Seq=64	Seq=128	Seq=256	Seq=512	Seq=1024
Softmax	4.5	6.0	7.4	8.6	9.8
QSoftmax	5.5	6.3	7.4	8.5	9.4
Linear	6.0	7.0	8.0	9.0	10.0

QSoftmax maintains slightly lower entropy than softmax at longer sequences.

D) Inference Time

Method	Seq=64	Seq=128	Seq=256	Seq=512	Seq=1024
Linear	0.07ms	0.11ms	0.29ms	0.75ms	2.3ms
Softmax	0.04ms	0.08ms	0.50ms	1.6ms	14ms
QSoftmax	14ms	57ms	190ms	530ms	1600ms

Note: Current Python implementation is not optimized. The C/NEON implementation provides significant speedups. Further optimization with vectorized complex arithmetic is planned.

Installation

Requirements

Python 3.8+
NumPy
Matplotlib (for visualization)
ARM processor (for NEON optimization)

Install

# Clone the repository
git clone https://github.com/yourusername/qsoftmax.git
cd qsoftmax

# Install Python dependencies
pip install numpy matplotlib

# Build C/NEON benchmark (optional)
make

Usage

Python

from qsoftmax import QSoftmax, QSoftmaxConfig

# Initialize with configuration
config = QSoftmaxConfig(
    num_basis_states=4,           # Quantum state dimension
    entanglement_strength=0.3,    # Entanglement coupling
    interference_threshold=0.1,   # Sparsity threshold
    mode="full"                   # "standard", "sparse", "entangled", "full"
)

qsoftmax = QSoftmax(config)

# Forward pass (drop-in replacement for softmax attention)
# Q, K, V: [seq_len, head_dim]
output, attention_weights = qsoftmax(Q, K, V)

Benchmarking

# Run full benchmark suite with graphs
python3 benchmark.py --output-dir benchmark_results

# Run Llama-2 style benchmark
python3 llama2_benchmark.py --output-dir llama2_results --int4

# Run C/NEON benchmark
make test

Docker

# Build container
docker build -t qsoftmax .

# Run benchmarks
docker run --rm -v $(pwd)/results:/app/benchmark_results qsoftmax

API Reference

QSoftmaxConfig

@dataclass
class QSoftmaxConfig:
    num_basis_states: int = 4
    # Number of quantum basis states (4 = 2-qubit, 8 = 3-qubit)

    entanglement_strength: float = 0.3
    # Coupling strength for key entanglement [0, 1]

    interference_threshold: float = 0.1
    # Scores below max * threshold are zeroed [0, 1]

    use_sparse: bool = True
    # Enable sparsity through thresholding

    mode: str = "full"
    # "standard": Basic QSoftmax
    # "sparse": With sparsity threshold
    # "entangled": With key entanglement
    # "full": All features enabled

QSoftmax.forward()

def forward(
    self,
    Q: np.ndarray,  # [seq_len, head_dim] queries
    K: np.ndarray,  # [seq_len, head_dim] keys
    V: np.ndarray   # [seq_len, head_dim] values
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Returns:
        output: [seq_len, head_dim] - Attention output
        attn_weights: [seq_len, seq_len] - Attention weight matrix
    """

Theoretical Background

Quantum Mechanics Analogy

Quantum Concept	QSoftmax Implementation
Superposition	Token encoded in multiple basis states
Amplitude	Complex number (magnitude + phase)
Interference	Phases add/cancel to compute scores
Entanglement	Correlated key states
Measurement	\|⟨q\|k⟩\|² gives attention probability

Why Interference Avoids Dilution (Theory)

In standard softmax:

P(best) = exp(max) / Σⱼ exp(sⱼ)
        → 0 as L → ∞

In QSoftmax:

P(best) = |⟨query|key_best⟩|²

Non-matching keys: phases misalign → |⟨q|k⟩|² → 0
Matching keys: phases align → |⟨q|k⟩|² → 1

No division by L → No dilution

Complex Matrix Multiplication

QSoftmax requires complex arithmetic. We implement efficient methods:

Method	Operations	Use Case
Naive	4 real matmuls	Baseline
Karatsuba	3 real matmuls	25% fewer ops
NEON	Vectorized complex	ARM optimization

Limitations & Future Work

Current Limitations

Needle Attention: Current threshold settings may be too aggressive, causing relevant information to be filtered. Requires careful tuning of interference_threshold.
Computational Cost: Python implementation is ~100x slower than optimized softmax. C/NEON implementation improves this significantly.
Encoding Sensitivity: The mapping from real vectors to quantum states affects performance. Alternative encodings (learned phases) may improve results.

Planned Improvements

Learnable phase encoding (end-to-end training)
Adaptive threshold based on score distribution
CUDA implementation for GPU acceleration
Integration with PyTorch/JAX autograd
Hybrid attention (softmax for local, QSoftmax for global)
Implement on SVE2, as underrepresented and can benefit from this framework.

Research Directions

Learned Interference: Train the encoding to maximize task performance
Multi-Qubit Extension: 8D or 16D states for more expressiveness
Hardware Acceleration: Custom FPGA/ASIC for complex arithmetic

File Structure

qsoftmax/
├── qsoftmax_types.h      # C type definitions
├── qsoftmax_neon.h       # NEON-optimized C implementation
├── benchmark_c.c         # C benchmark
├── qsoftmax.py           # Python implementation
├── benchmark.py          # Python benchmark + visualization
├── llama2_benchmark.py   # Llama-2 style benchmark
├── Makefile              # Build system
├── Dockerfile            # ARM64 container
└── README.md             # This file

Citation

If you use QSoftmax in your research, please cite:

@software{qsoftmax2024,
  title={QSoftmax: Quantum-Inspired Softmax for Attention Mechanisms},
  author={Aksh Parekh},
  year={2026},
  url={https://github.com/aparekh02/qsoftmax}
}

Related Work

Attention Is All You Need - Original Transformer
Quantum Decision Transformers - Quantum-inspired RL
Sparsemax - Sparse attention alternative
Linear Attention - O(L) complexity attention

License

This project is licensed under the MIT License.

Acknowledgments

ARM NEON intrinsics documentation
NumPy and Matplotlib communities
Quantum computing education resources

QSoftmax: Where quantum mechanics meets attention mechanisms.

Name		Name	Last commit message	Last commit date
benchmark_results		benchmark_results
qsoftmax		qsoftmax
.DS_Store		.DS_Store
CBenchmark_results.txt		CBenchmark_results.txt
Dockerfile		Dockerfile
README.md		README.md
docker-compose.yml		docker-compose.yml
qsoftmax_onnx.py		qsoftmax_onnx.py
requirements-arm.txt		requirements-arm.txt
run_arm_docker.sh		run_arm_docker.sh

Latest commit

History

Folders and files

Repository files navigation