85 lines
4.4 KiB
Markdown
85 lines
4.4 KiB
Markdown
|
---
|
|||
|
|
license: mit
|
|||
|
|
pipeline_tag: text-generation
|
|||
|
|
tags: [research, experimental, gravity-attention, qwen2]
|
|||
|
|
---
|
|||
|
|
|
|||
|
|
# Gravity-2
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
**Experimental research model by squ11z1.**
|
|||
|
|
|
|||
|
|
A 3B reasoning model in which the standard
|
|||
|
|
scaled-dot-product attention is replaced by a physically-motivated **gravity attention**,
|
|||
|
|
then adapted with LoRA. This card documents a **stage-1 proof-of-mechanism**
|
|||
|
|
|
|||
|
|
## The experiment
|
|||
|
|
|
|||
|
|
Transformer attention scores tokens by **alignment** — the dot product `q·k`. Gravity-2
|
|||
|
|
asks a different question: *what if tokens attended by **proximity** instead?* We replace
|
|||
|
|
the score with an inverse-square law borrowed from gravitation — each token is pulled
|
|||
|
|
toward others that are close in query/key space, weighted by a learnable per-head "mass":
|
|||
|
|
|
|||
|
|
```
|
|||
|
|
M_h²
|
|||
|
|
score(i, j) = ───────────────────── → softmax_j( score )
|
|||
|
|
‖q_i − k_j‖² + ε
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
- **M_h = softplus(gravity_mass_log[h])** — one learnable mass per **query head** (16 / layer),
|
|||
|
|
initialised at 0.5; `softplus` keeps it strictly positive.
|
|||
|
|
- **‖q_i − k_j‖²** — squared L2 distance, computed stably as `‖q‖² + ‖k‖² − 2·q·k`.
|
|||
|
|
- **ε = 0.1** — softening length; prevents the `q → k` singularity.
|
|||
|
|
- The raw gravity scores are then passed through the **usual softmax** (see Limitations).
|
|||
|
|
|
|||
|
|
### Why it's interesting
|
|||
|
|
- **Different inductive bias.** Dot-product attention rewards directional alignment;
|
|||
|
|
inverse-distance rewards *locality* in the learned embedding geometry — a metric prior
|
|||
|
|
rather than an inner-product one.
|
|||
|
|
- **Interpretable per-head masses.** Each head learns a scalar "mass" controlling how
|
|||
|
|
sharply it concentrates — a compact, inspectable knob (see `figures/04_mass_heatmap.png`).
|
|||
|
|
- **A bridge to physics-style sparsity.** An inverse-square field is naturally local, which
|
|||
|
|
later stages (pruning / QUBO, "Gravity-6") aim to exploit for structured sparsity.
|
|||
|
|
|
|||
|
|
## Architecture
|
|||
|
|
Qwen2-3B class: 36 layers, hidden 2048, **16 query heads / 2 KV heads (GQA, group size 8)**,
|
|||
|
|
head_dim 128. The 2 KV heads are `repeat_kv`-expanded to 16 before the distance, so each
|
|||
|
|
query head gets its own mass. Integrated via the transformers-5.x `AttentionInterface`
|
|||
|
|
(a registered `"gravity"` op + eager causal-mask reuse) — RoPE / KV-cache / masking are
|
|||
|
|
left to the framework; only the score function changes.
|
|||
|
|
|
|||
|
|
## Results
|
|||
|
|
|
|||
|
|
| | |
|
|||
|
|
|---|---|
|
|||
|
|
|  |  |
|
|||
|
|
|  |  |
|
|||
|
|
|  |  |
|
|||
|
|
|
|||
|
|
## Honest limitations
|
|||
|
|
- **Not "pure" gravity.** The inverse-square scores are renormalised by a **softmax on top**
|
|||
|
|
(`softmax_j(M²/(d²+ε))`). Without it training was unstable, but it means this is a
|
|||
|
|
*distance-biased softmax attention*, not a literal gravitational field — the normalisation
|
|||
|
|
reintroduces global competition between keys.
|
|||
|
|
- **MHA → GQA transfer is an open question.** The mechanism was first prototyped on MHA
|
|||
|
|
(1 KV head per query head). Here it runs on GQA by `repeat_kv`-expanding 2 KV heads to 16
|
|||
|
|
and giving each query head its own mass; whether this is the right granularity (vs. one
|
|||
|
|
mass per KV group) is **unresolved** and may matter for convergence.
|
|||
|
|
- **Loading requires the patch** (below). **GGUF builds run standard attention, not gravity**
|
|||
|
|
(llama.cpp has no kernel for `M²/(‖q−k‖²+ε)`) — the `*.gguf` files are format placeholders
|
|||
|
|
and produce incorrect output.
|
|||
|
|
|
|||
|
|
## Loading (requires the gravity patch)
|
|||
|
|
```bash
|
|||
|
|
python load_gravity2.py # from_pretrained -> patch_qwen_with_gravity -> load gravity_mass_log.pt
|
|||
|
|
```
|
|||
|
|
Weights are LoRA-merged into the base but were trained under gravity scoring; loading them
|
|||
|
|
under vanilla attention gives garbage. `config.json` ships `_attn_implementation="eager"`
|
|||
|
|
only so the checkpoint loads — the patch switches it to gravity.
|
|||
|
|
|
|||
|
|
## License & attribution
|
|||
|
|
Released under the **MIT License**. This is a **derivative work of
|
|||
|
|
[`WeiboAI/VibeThinker-3B`](https://huggingface.co/WeiboAI/VibeThinker-3B)** (the base model
|
|||
|
|
for the experiment), which is distributed under the **MIT License**; that license is
|
|||
|
|
inherited here and the original authors are credited accordingly.
|