QKVE-61M: Learned Value Gating in Transformer Attention

English · Polski

---

English

Overview

QKVE-61M is an experimental compact language model (61.03M parameters) designed to explore learned value modulation within Transformer attention mechanisms.

The project originated as an empirical, hands-on journey to deeply understand how Transformers work under the hood. Specifically, it investigated whether horizontal expressiveness could be augmented through dedicated gating projections rather than merely stacking deeper blocks. Through iterative experimentation on a single consumer NVIDIA GeForce RTX 3090 (24 GB VRAM), this exploration evolved into an efficient and agile research platform for training, ablation, and architecture probing.

The concrete result of this research platform is the QKVE-61M model released here.

Architecture & Mathematical Formulation

In standard Transformer attention, value representations are directly aggregated using softmax attention weights: V(x) = x * W_V.

In contrast, the QKVE architecture introduces an additional learned projection matrix W_E, forming an input-dependent dynamic modulation gate over the value vectors before attention aggregation. The letter E in QKVE derives from Energy (extending standard attention to Query, Key, Value, Energy), conceptualizing an input-dependent gating of activation energy across value channels:

$$V_{\text{mod}}(x)=(x{W}_V)\odot (1+\tanh (x{W}_E))V\_\{\backslash text\{mod\}\}(x)\; =\; (x\; W\_V)\; \backslash odot\; \backslash left(1\; +\; \backslash tanh(x\; W\_E)\backslash right)$$

Component Breakdown

Component	Tensor Shape / Range	Role in Attention
x	[Batch, Seq_Len, 384]	Normalized input representation
W_V	[384, 128]	Learned value projection matrix (4 heads × 32 dim)
W_E	[384, 128]	Learned value-gate projection matrix ("E" for Energy, 4 heads × 32 dim)
⊙	Elementwise	Hadamard product (elementwise multiplication)
1 + tanh(x * W_E)	Range: (0, 2)	Value modulation multiplier (attenuates or amplifies without sign flip)

• Origin of "E" (Energy): The letter E in QKVE stands for Energy — a learned, state-dependent gating mechanism that scales the energetic intensity of value components based on input context.

• No Sign Inversion: Standard tanh outputs span (-1, 1), which could invert feature signs. The 1 + tanh(x * W_E) formulation shifts the multiplier strictly to positive values in (0, 2). A gate output of 1.0 acts as neutral pass-through (preserving standard attention behavior), while values < 1.0 attenuate noise and values > 1.0 amplify salient features.

• Grouped Query Attention (GQA): The gate W_E matches the Grouped Query Attention structure (4 KV/Gate heads for 12 Query heads, 3:1 ratio), modulating grouped values before causal attention aggregation.

Model Specifications

Parameter	Value
Total Parameters	61,030,016 (~61M)
Layers / Blocks	30 (QKVE30Block)
Hidden Dimension (d_model)	384
Attention Heads (Query)	12 (head dim = 32)
KV & Gate Heads (num_kv_heads)	4 (Grouped Query Attention, GQA 3:1)
Feed-Forward Network	Deep SwiGLU (2 sub-layers per block = 60 FFN layers, d_ffn = 512)
Positional Embeddings	Rotary Position Embeddings (RoPE)
Normalization	RMSNorm (eps = 1e-6)
Vocabulary Size	32,128 tokens (SpeakLeash Bielik-11B tokenizer)
Context Length	512 tokens (up to 1024)
Weight Tying	lm_head tied to tok_embeddings

Training Setup & Platform

• Hardware: 1× NVIDIA GeForce RTX 3090 (24 GB VRAM)
• Pretraining Budget: ~1.08 Billion tokens across curated Polish and English corpora.
• Fine-Tuning: Supervised Fine-Tuning (SFT) followed by Chat SFT (ChatML format).
• Final Metrics: Training Loss: ~1.48, Perplexity: ~4.40 on conversational data.
• Efficiency: Achieved throughput of ~9,800 tokens/second on a single consumer GPU while reserving only ~11.5 GB VRAM.

Available Formats

1. model.safetensors: SafeTensors format for PyTorch / Hugging Face ecosystem.
2. pytorch_model.bin: Standard PyTorch state dict.
3. qkve-61m.pt: Exported TorchScript archive for C++ / mobile / zero-dependency deployment.
4. qkve-61m.onnx & qkve-61m-int8.onnx: Full-precision and INT8 quantized ONNX models (ready for ONNX Runtime / Android Pixel 6a).

Quickstart (Python)

1. Install Dependencies

pip install torch transformers safetensors huggingface_hub

2. Complete Copy-and-Run Inference Example

This self-contained script downloads the model architecture (model.py), weights (model.safetensors), and tokenizer from Hugging Face, strictly verifies the loaded weights, and generates text:

import os
import sys
import torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
from safetensors.torch import load_file

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# 1. Download model files from Hugging Face
model_dir = snapshot_download(repo_id="Sticzu/QKVE-61M")
sys.path.insert(0, model_dir)
from model import QKVE30Transformer

# 2. Load tokenizer and instantiate architecture
tokenizer = AutoTokenizer.from_pretrained(model_dir, clean_up_tokenization_spaces=False)
model = QKVE30Transformer(
    vocab_size=len(tokenizer),
    d_model=384,
    num_layers=30,
    num_heads=12,
    num_kv_heads=4,
    base_ffn=512,
    max_seq_len=512
).to(device)

# 3. Load and strictly verify weights
weights = load_file(os.path.join(model_dir, "model.safetensors"))
load_result = model.load_state_dict(weights, strict=False)
if load_result.missing_keys:
    raise RuntimeError(f"Missing keys: {load_result.missing_keys}")
if load_result.unexpected_keys:
    raise RuntimeError(f"Unexpected keys: {load_result.unexpected_keys}")

model.eval()

# 4. Generate text (ChatML format)
prompt = "<|im_start|>user\nWytłumacz w jednym zdaniu czym jest atencja w sztucznej inteligencji.<|im_end|>\n<|im_start|>assistant\n"
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
eos_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")

with torch.inference_mode():
    output_ids = model.generate(
        input_ids,
        max_new_tokens=80,
        temperature=0.7,
        top_p=0.9,
        eos_token_id=eos_token_id,
    )

answer_ids = output_ids[0, input_ids.shape[1]:]
print(tokenizer.decode(answer_ids, skip_special_tokens=True).strip())

Experimental Scope & Proof-of-Concept Note

Experimental Validation:
This model is an experimental research prototype created to empirically verify whether the proposed QKVE architecture—specifically incorporating input-dependent value modulation via a learned Energy gate $W_E$ alongside Grouped-Query Attention (GQA) and deep dual SwiGLU FFNs—can be trained from scratch, maintain numerical stability across 30 layers, converge efficiently, and generate coherent natural language.

At 61.03M parameters, this model serves as a technical proof-of-concept (PoC) demonstrating architecture viability rather than an instruction-tuned assistant or calculator. It exhibits standard characteristics of compact pre-trained language models, including occasional factual drift and lack of grounded arithmetic reasoning.

Sample Generations (RTX 3090 Test Runs)

The following generations were produced directly by the model on an NVIDIA RTX 3090 using the ChatML template (<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n, temperature=0.6, top_p=0.9):

1. Identity & Role

• Prompt: Kim jesteś i w czym możesz pomóc? (Who are you and how can you help?)
• Response:

  Jestem modelem języka AI i nie mam możliwości wykonywania zadań, aby pomóc. Witam, czy możesz mi powiedzieć, w jaki sposób możesz pomóc?

• Observation: The model correctly identifies itself as an AI language model and responds in coherent, grammatically sound Polish.

2. Technical Domain Understanding

• Prompt: Czym jest mechanizm atencji w modelach językowych? (What is the attention mechanism in language models?)
• Response:

  W modelach językowych, mechanizm atencji jest wykorzystywany do analizy zachowań i interakcji między różnymi elementami, takimi jak wartości, struktury i reprezentacje...

• Observation: Successfully activates domain-specific representations (attention, interactions, values, structures) aligned with transformer concepts.

3. Factual Recall & Polish General Knowledge

• Prompt: Stolicą Polski jest: (The capital of Poland is:)
• Response:

  Stolicą Polski jest Warszawa.

• Observation: Accurate factual recall from the Polish SpeakLeash pretraining corpora.

4. Arithmetic & Reasoning Probing (Small-Model Behavior)

• Prompt: Ile to jest 2 + 2? (How much is 2 + 2?)
• Response:

  2 + 2 to 3,5.

• Observation: Demonstrates classic compact LLM token association without external execution tools: the model generates statistically plausible numerical sequences rather than executing symbolic arithmetic.

---

Polski

Wprowadzenie

QKVE-61M to eksperymentalny, lekki model językowy (61,03M parametrów) stworzony w celu zbadania mechanizmu uczonego bramkowania wektorów wartości w mechanizmie uwagi (Attention).

Projekt rozpoczął się jako eksperymentalna, praktyczna próba dogłębnego zrozumienia działania architektury Transformer od podstaw. Głównym założeniem badawczym było pytanie: czy możemy zwiększyć pojemność ekspresyjną mechanizmu uwagi horyzontalnie — poprzez dodanie dedykowanej projekcji bramkującej modułu atencji — zamiast standardowego zwiększania liczby warstw?

Dzięki intensywnym iteracjom na pojedynczej konsumenckiej karcie graficznej NVIDIA GeForce RTX 3090 (24 GB VRAM), projekt przekształcił się w niezwykle wydajną i elastyczną platformę do eksperymentowania z dynamiką treningu oraz architekturą transformera.

Praktycznym zwieńczeniem i dowodem działania tej platformy badawczej jest udostępniony tutaj model QKVE-61M.

Wzór matematyczny i mechanizm działania

W klasycznym module Attention wektory wartości są łączone przy użyciu wag z rozkładu softmax: V(x) = x * W_V.

Architektura QKVE wprowadza dodatkową, wyuczoną macierz projekcji W_E, która generuje dynamiczną bramkę modulującą wektory wartości przed wyznaczeniem wag atencji. Litera E w nazwie QKVE wywodzi się bezpośrednio od słowa Energy (Energia) (rozszerzenie klasycznego układu do Query, Key, Value, Energy), odzwierciedlając wyuczoną modulację intensywności i przepływu energii w kanałach wartości:

$$V_{\text{mod}}(x)=(x{W}_V)\odot (1+\tanh (x{W}_E))V\_\{\backslash text\{mod\}\}(x)\; =\; (x\; W\_V)\; \backslash odot\; \backslash left(1\; +\; \backslash tanh(x\; W\_E)\backslash right)$$

Zestawienie komponentów i wymiary

Komponent	Kształt tensora / Zakres	Rola w module Attention
x	[Batch, Seq_Len, 384]	Znormalizowana reprezentacja wejściowa
W_V	[384, 128]	Wyuczona macierz projekcji wartości (4 głowy × 32 dim)
W_E	[384, 128]	Wyuczona macierz projekcji bramki wartości („E” od Energy / Energia, 4 głowy × 32 dim)
⊙	Po współrzędnych	Iloczyn Hadamarda (mnożenie wektorów element po elemencie)
1 + tanh(x * W_E)	Zakres: (0, 2)	Mnożnik modulacji (płynne tłumienie lub wzmacnianie bez zmiany znaku)

• Geneza litery „E” (Energy): Skrót QKVE rozszerza klasyczną trójkę QKV o komponent Energy (Energia) — wyuczoną, zależną od kontekstu bramkę modulującą poziom energii i istotności wektorów wartości dla poszczególnych tokenów.

• Brak odwracania znaku: Standardowa funkcja tanh przyjmuje wartości w przedziale (-1, 1), co mogłoby odwrócić znak składowych cech. Formuła 1 + tanh(x * W_E) przesuwa ten zakres ściśle do wartości dodatnich w przedziale (0, 2). Wartość 1.0 odpowiada neutralnemu przepuszczeniu cech (jak w standardowym attention), wartości < 1.0 tłumią szum, a wartości > 1.0 selektywnie wzmacniają istotne informacje.

• Grouped Query Attention (GQA): Projekcja bramki W_E ma identyczny kształt jak projekcja wartości W_V (4 głowy KV i Gate dla 12 głów Query, relacja 3:1), modulując grupy wartości przed operacją causal scaled dot-product attention.

Specyfikacja architektury

Cecha	Wartość
Liczba parametrów	61 030 016 (~61M)
Liczba bloków	30 bloków (QKVE30Block)
Wymiar ukryty (d_model)	384
Głowy uwagi Query	12 głów (wymiar głowy = 32)
Głowy KV i Bramki E (num_kv_heads)	4 głowy (Grouped Query Attention, relacja 3:1)
Warstwy FFN	Głęboki SwiGLU (2 podwarstwy na blok = łącznie 60 warstw FFN, d_ffn = 512)
Pozycjonowanie	Rotary Position Embeddings (RoPE)
Normalizacja	RMSNorm (eps = 1e-6)
Tokenizator	32 128 tokenów (SpeakLeash Bielik-11B)
Maksymalny kontekst	512 tokenów (obsługa do 1024)
Wiązanie wag	lm_head współdzieli wagi z embeddingami wejściowymi

Przebieg treningu i wydajność

• Sprzęt treningowy: 1× NVIDIA GeForce RTX 3090 (24 GB VRAM).
• Wolumen treningowy: Ponad 1,08 miliarda tokenów (pretraining na zbalansowanych korpusach polskich i angielskich).
• Etapy strojenia: Pretrening (3 etapy) → Supervised Fine-Tuning (SFT) → Chat SFT (format ChatML).
• Końcowe metryki: Strata (loss): ~1.48, perplexity: ~4.40.
• Wydajność: Trening osiągał przepustowość rzędu ~9 800 tokenów/s przy zużyciu jedynie ~11,5 GB VRAM, co czyni model idealną bazą do dalszych badań i wdrożeń brzegowych (np. na smartfonach).

Dostępne formaty

• model.safetensors: Nowoczesny, bezpieczny format wag dla PyTorch.
• pytorch_model.bin: Klasyczny zrzut słownika wag (state_dict).
• qkve-61m.pt: Zoptymalizowany model w formacie TorchScript (możliwość uruchomienia bez kodu Pythona / na urządzeniach mobilnych).
• qkve-61m.onnx oraz qkve-61m-int8.onnx: Modele w formacie ONNX (w tym skwantyzowany INT8 zajmujący zaledwie ~64 MB).
• inference.py: Gotowy, przetestowany skrypt do prowadzenia konwersacji w terminalu.

Szybki start (Python)

1. Instalacja zależności

pip install torch transformers safetensors huggingface_hub

2. Kompletny skrypt uruchomieniowy „wklej i uruchom”

Poniższy skrypt automatycznie pobiera architekturę (model.py), wagi (model.safetensors) i tokenizator z Hugging Face przez snapshot_download, weryfikuje poprawność wag i generuje odpowiedź:

import os
import sys
import torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
from safetensors.torch import load_file

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Używane urządzenie: {device}")

# 1. Pobranie plików modelu z Hugging Face
model_dir = snapshot_download(repo_id="Sticzu/QKVE-61M")
sys.path.insert(0, model_dir)
from model import QKVE30Transformer

# 2. Załadowanie tokenizatora i inicjalizacja architektury
tokenizer = AutoTokenizer.from_pretrained(model_dir, clean_up_tokenization_spaces=False)
model = QKVE30Transformer(
    vocab_size=len(tokenizer),
    d_model=384,
    num_layers=30,
    num_heads=12,
    num_kv_heads=4,
    base_ffn=512,
    max_seq_len=512
).to(device)

# 3. Kontrola poprawności i ładowanie wag
weights = load_file(os.path.join(model_dir, "model.safetensors"))
load_result = model.load_state_dict(weights, strict=False)
if load_result.missing_keys:
    raise RuntimeError(f"Brakujące wagi: {load_result.missing_keys}")
if load_result.unexpected_keys:
    raise RuntimeError(f"Nieoczekiwane klucze: {load_result.unexpected_keys}")

model.eval()

# 4. Generowanie odpowiedzi (format ChatML)
prompt = "<|im_start|>user\nWytłumacz w jednym zdaniu, czym jest atencja w sieciach neuronowych.<|im_end|>\n<|im_start|>assistant\n"
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
eos_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")

with torch.inference_mode():
    output_ids = model.generate(
        input_ids,
        max_new_tokens=80,
        temperature=0.7,
        top_p=0.9,
        eos_token_id=eos_token_id,
    )

answer_ids = output_ids[0, input_ids.shape[1]:]
print(tokenizer.decode(answer_ids, skip_special_tokens=True).strip())

---

Eksperymentalny charakter i Proof-of-Concept

Charakter badawczo-testowy:
Udostępniony model QKVE-61M powstał jako eksperymentalny prototyp badawczy (Proof-of-Concept), którego nadrzędnym celem było empiryczne zweryfikowanie, czy nowo zaproponowana architektura z nieliniowym bramkowaniem wartości (bramka energii $W_E$ w mechanizmie uwagi) w połączeniu z GQA i podwójnym SwiGLU jest w stanie stabilnie się uczyć od zera, zbiegać numerycznie w głębokiej 30-warstwowej sieci i opanować modelowanie języka naturalnego.

Przy skali 61,03M parametrów model stanowi techniczny dowód koncepcji, a nie docelowy produkt komercyjny czy asystenta obliczeniowego. Posiada on cechy typowe dla małych modeli bazowych: poprawnie operuje gramatyką i podstawowymi faktami, lecz wykazuje halucynacje przy ścisłych operacjach logiczno-matematycznych bez zewnętrznych narzędzi.

Przykładowe odpowiedzi modelu (RTX 3090)

Poniższe odpowiedzi zostały wygenerowane bezpośrednio przez model na karcie NVIDIA RTX 3090 z użyciem szablonu ChatML (<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n, temperature=0.6, top_p=0.9):

1. Tożsamość i autoprezentacja

• Pytanie: Kim jesteś i w czym możesz pomóc?
• Odpowiedź:

  Jestem modelem języka AI i nie mam możliwości wykonywania zadań, aby pomóc. Witam, czy możesz mi powiedzieć, w jaki sposób możesz pomóc?

• Komentarz: Model bezbłędnie identyfikuje swoją rolę jako sztuczna inteligencja i formułuje poprawną składniowo wypowiedź po polsku.

2. Zrozumienie domeny technicznej

• Pytanie: Czym jest mechanizm atencji w modelach językowych?
• Odpowiedź:

  W modelach językowych, mechanizm atencji jest wykorzystywany do analizy zachowań i interakcji między różnymi elementami, takimi jak wartości, struktury i reprezentacje...

• Komentarz: Model poprawnie operuje terminologią atencji, interakcji i struktur reprezentacji w sieciach neuronowych.

3. Podstawowa wiedza ogólna

• Pytanie: Stolicą Polski jest:
• Odpowiedź:

  Stolicą Polski jest Warszawa.

• Komentarz: Prawidłowe przywołanie faktów z polskiego korpusu SpeakLeash.

4. Test rozumowania arytmetycznego (ograniczenia małej skali)

• Pytanie: Ile to jest 2 + 2?
• Odpowiedź:

  2 + 2 to 3,5.

• Komentarz: Doskonała ilustracja statystycznej natury małych modeli językowych (~61M). Bez dedykowanego kalkulatora lub finetuningu arytmetycznego model traktuje cyfry jako tokeny językowe.

---

Licencja i powiązania

Kod modelu i wagi udostępniane są na licencji Apache 2.0. Repozytorium kodu badawczego: GitHub: Sticzu/qkve-transforer-architecture.