commit 20ac911fcf74ab00d2722e7e5a6ebbaffc1e9a34 Author: ModelHub XC Date: Tue Aug 11 07:47:17 2026 +0800 初始化项目，由ModelHub XC社区提供模型 Model: squ11z1/Gravity-2 Source: Original Platform diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..bb82cc7 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,38 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text +tokenizer.json filter=lfs diff=lfs merge=lfs -text +gravity-2-f16.gguf filter=lfs diff=lfs merge=lfs -text +gravity-2-Q4_K_M.gguf filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md new file mode 100644 index 0000000..9234a54 --- /dev/null +++ b/README.md @@ -0,0 +1,85 @@ +--- +license: mit +pipeline_tag: text-generation +tags: [research, experimental, gravity-attention, qwen2] +--- + +# Gravity-2 + +![IMAGE 2026-06-16 19:46:27](https://cdn-uploads.huggingface.co/production/uploads/67329d3f69fded92d56ab41a/NkBrWQXsuwZUrsTaYga5-.jpeg) + +**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 + +| | | +|---|---| +| ![loss](figures/01_loss.png) | ![masses](figures/02_mass_band.png) | +| ![grad](figures/03_gradnorm.png) | ![heatmap](figures/04_mass_heatmap.png) | +| ![aer](figures/05_aer_entropy.png) | ![concept](figures/06_concept.png) | + +## 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. \ No newline at end of file diff --git a/chat_template.jinja b/chat_template.jinja new file mode 100644 index 0000000..28028c0 --- /dev/null +++ b/chat_template.jinja @@ -0,0 +1,54 @@ +{%- if tools %} + {{- '<|im_start|>system\n' }} + {%- if messages[0]['role'] == 'system' %} + {{- messages[0]['content'] }} + {%- else %} + {{- 'You are a helpful assistant.' }} + {%- endif %} + {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} +{%- else %} + {%- if messages[0]['role'] == 'system' %} + {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }} + {%- else %} + {{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- for message in messages %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %} + {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {{- '<|im_start|>' + message.role }} + {%- if message.content %} + {{- '\n' + message.content }} + {%- endif %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {{- '\n\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {{- tool_call.arguments | tojson }} + {{- '}\n' }} + {%- endfor %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- message.content }} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} +{%- endif %} diff --git a/config.json b/config.json new file mode 100644 index 0000000..67ee208 --- /dev/null +++ b/config.json @@ -0,0 +1,69 @@ +{ + "architectures": [ + "Qwen2ForCausalLM" + ], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "dtype": "bfloat16", + "eos_token_id": 151643, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 11008, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 131072, + "max_window_layers": 36, + "model_type": "qwen2", + "num_attention_heads": 16, + "num_hidden_layers": 36, + "num_key_value_heads": 2, + "pad_token_id": null, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "rope_theta": 1000000.0, + "rope_type": "default" + }, + "sliding_window": null, + "tie_word_embeddings": true, + "transformers_version": "5.12.1", + "use_cache": false, + "use_sliding_window": false, + "vocab_size": 151936 +} diff --git a/figures/01_loss.png b/figures/01_loss.png new file mode 100644 index 0000000..1e1781d Binary files /dev/null and b/figures/01_loss.png differ diff --git a/figures/02_mass_band.png b/figures/02_mass_band.png new file mode 100644 index 0000000..bb093c4 Binary files /dev/null and b/figures/02_mass_band.png differ diff --git a/figures/03_gradnorm.png b/figures/03_gradnorm.png new file mode 100644 index 0000000..fdfecf4 Binary files /dev/null and b/figures/03_gradnorm.png differ diff --git a/figures/04_mass_heatmap.png b/figures/04_mass_heatmap.png new file mode 100644 index 0000000..f62199d Binary files /dev/null and b/figures/04_mass_heatmap.png differ diff --git a/figures/05_aer_entropy.png b/figures/05_aer_entropy.png new file mode 100644 index 0000000..e3de38c Binary files /dev/null and b/figures/05_aer_entropy.png differ diff --git a/figures/06_concept.png b/figures/06_concept.png new file mode 100644 index 0000000..eef389d Binary files /dev/null and b/figures/06_concept.png differ diff --git a/generation_config.json b/generation_config.json new file mode 100644 index 0000000..cf72f87 --- /dev/null +++ b/generation_config.json @@ -0,0 +1,6 @@ +{ + "bos_token_id": 151643, + "eos_token_id": 151643, + "max_new_tokens": 2048, + "transformers_version": "5.12.1" +} diff --git a/gravity-2-Q4_K_M.gguf b/gravity-2-Q4_K_M.gguf new file mode 100644 index 0000000..f9c2921 --- /dev/null +++ b/gravity-2-Q4_K_M.gguf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:68492d7cc67d8796ff80a06347a8f7c7754d164a711f531bce4dfedfc4b646f7 +size 1929902304 diff --git a/gravity-2-f16.gguf b/gravity-2-f16.gguf new file mode 100644 index 0000000..295d1e3 --- /dev/null +++ b/gravity-2-f16.gguf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b1cab036db87b080b897cfc4637219bece09c9bee7634c9bce5009abcf5d0b53 +size 6178316512 diff --git a/gravity_attention_qwen.py b/gravity_attention_qwen.py new file mode 100644 index 0000000..7af1128 --- /dev/null +++ b/gravity_attention_qwen.py @@ -0,0 +1,110 @@ +""" +Gravity-2 attention for Qwen2 / VibeThinker-3B (transformers 5.x interface). + +Replaces softmax(QKᵀ·scaling) with a physically-motivated score: + + score(i,j) = M_h² / (||q_i − k_j||² + eps) # then standard softmax over j + + • M_h = softplus(gravity_mass_log[h]) — one learnable mass per QUERY head (16/layer) + • ||q_i − k_j||² = ||q||² + ||k||² − 2·q·k # GQA: K repeated 2→16 first + • eps guards the singularity at q==k + +Integration uses the transformers-5.x AttentionInterface dispatch (NOT a forward +monkeypatch): we register a "gravity" attention fn + alias its mask to "eager" so +the framework keeps building the additive causal mask, handling RoPE/cache itself. +""" +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.models.qwen2.modeling_qwen2 import repeat_kv +from transformers.modeling_utils import AttentionInterface +from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS + +ATTN_NAME = "gravity" + + +def gravity_attention_forward(module, query, key, value, attention_mask, + scaling=None, dropout=0.0, **kwargs): + """AttentionInterface contract. + + query: (B, Hq, Tq, D) key/value: (B, Hkv, Tk, D) + returns: (attn_output (B, Tq, Hq, D), attn_weights (B, Hq, Tq, Tk)) + `scaling` is intentionally ignored — gravity replaces the 1/√d scale. + """ + # GQA: expand 2 KV heads up to 16 so distances live in per-query-head space + key = repeat_kv(key, module.num_key_value_groups) + value = repeat_kv(value, module.num_key_value_groups) + + # ||q_i - k_j||^2 in fp32 for numerical stability + q = query.float() + k = key.float() + q_sq = (q * q).sum(-1, keepdim=True) # (B,Hq,Tq,1) + k_sq = (k * k).sum(-1, keepdim=True).transpose(-2, -1) # (B,Hq,1,Tk) + qk = torch.matmul(q, k.transpose(-2, -1)) # (B,Hq,Tq,Tk) + d_sq = (q_sq + k_sq - 2.0 * qk).clamp_min(0.0) + + mass = F.softplus(module.gravity_mass_log).float().view(1, -1, 1, 1) # (1,Hq,1,1) + scores = (mass * mass) / (d_sq + module.gravity_eps) # (B,Hq,Tq,Tk), fp32 + + if attention_mask is not None: + # additive causal mask (eager-style), already correct length + scores = scores + attention_mask[..., : key.shape[-2]].float() + + attn = F.softmax(scores, dim=-1, dtype=torch.float32) + + # AER: optionally stash mean per-row attention entropy (flag-gated, ~free when off) + if getattr(module, "_capture_entropy", False): + ent = -(attn.clamp_min(1e-12) * attn.clamp_min(1e-12).log()).sum(-1) + module._last_entropy = ent.mean().detach() + + attn = F.dropout(attn, p=dropout, training=module.training) + attn = attn.to(value.dtype) + + out = torch.matmul(attn, value) # (B,Hq,Tq,D) + out = out.transpose(1, 2).contiguous() # (B,Tq,Hq,D) + return out, attn + + +_REGISTERED = False + + +def _register(): + global _REGISTERED + if _REGISTERED: + return + AttentionInterface.register(ATTN_NAME, gravity_attention_forward) + # reuse the eager additive-mask builder for our custom impl + ALL_MASK_ATTENTION_FUNCTIONS.register(ATTN_NAME, ALL_MASK_ATTENTION_FUNCTIONS["eager"]) + _REGISTERED = True + + +def patch_qwen_with_gravity(model, eps: float = 0.1, init_mass: float = 0.5): + """Add per-head gravity_mass_log to every Qwen2 self-attn and switch dispatch. + + Leaves q/k/v/o_proj weights untouched. gravity_mass_log kept in fp32. + """ + _register() + init_log = math.log(math.exp(init_mass) - 1.0) # softplus^{-1}(init_mass) + H = model.config.num_attention_heads + n = 0 + for layer in model.model.layers: + attn = layer.self_attn + dev = attn.q_proj.weight.device + attn.gravity_mass_log = nn.Parameter( + torch.full((H,), init_log, device=dev, dtype=torch.float32) + ) + attn.gravity_eps = float(eps) + # config object is shared, but set defensively + attn.config._attn_implementation = ATTN_NAME + n += 1 + model.config._attn_implementation = ATTN_NAME + print(f"[gravity] patched {n} Qwen2 layers (heads={H}, eps={eps}, init_mass={init_mass})") + return model + + +def gravity_mass_state_dict(model): + """Extract only the gravity_mass_log params (for saving separately from base).""" + return {f"model.layers.{i}.self_attn.gravity_mass_log": + layer.self_attn.gravity_mass_log.detach().cpu() + for i, layer in enumerate(model.model.layers)} diff --git a/gravity_mass_log.pt b/gravity_mass_log.pt new file mode 100644 index 0000000..72b3c02 --- /dev/null +++ b/gravity_mass_log.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b37e7594676837dbf0b980970007938a992c059effcd5989c46e98ca1ced4c84 +size 14513 diff --git a/load_gravity2.py b/load_gravity2.py new file mode 100644 index 0000000..61d1f5f --- /dev/null +++ b/load_gravity2.py @@ -0,0 +1,16 @@ +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from gravity_attention_qwen import patch_qwen_with_gravity + +REPO = "." # or "squ11z1/Gravity-2" +tok = AutoTokenizer.from_pretrained(REPO) +model = AutoModelForCausalLM.from_pretrained(REPO, dtype=torch.bfloat16, + device_map="cuda", attn_implementation="eager") +patch_qwen_with_gravity(model) # re-enable gravity attention +masses = torch.load(f"{REPO}/gravity_mass_log.pt", map_location="cuda") +for i, layer in enumerate(model.model.layers): + layer.self_attn.gravity_mass_log.data.copy_(masses[f"model.layers.{i}.self_attn.gravity_mass_log"].cuda()) +model.eval() +ids = tok.apply_chat_template([{"role":"user","content":"What is 24*17?"}], + add_generation_prompt=True, return_tensors="pt", return_dict=True)["input_ids"].cuda() +print(tok.decode(model.generate(ids, max_new_tokens=200)[0, ids.shape[1]:], skip_special_tokens=True)) diff --git a/model.safetensors b/model.safetensors new file mode 100644 index 0000000..768dad5 --- /dev/null +++ b/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e645da04da9ff8321fa0f2f9894b7db9975d416ecd55e70e03af4dde88ccdfac +size 6171933008 diff --git a/tokenizer.json b/tokenizer.json new file mode 100644 index 0000000..55e03a0 --- /dev/null +++ b/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b17e16899b7fab7e695509f84bac5f10ed12804f0a590e935941e5af7f092f7f +size 11422263 diff --git a/tokenizer_config.json b/tokenizer_config.json new file mode 100644 index 0000000..493c153 --- /dev/null +++ b/tokenizer_config.json @@ -0,0 +1,16 @@ +{ + "add_prefix_space": false, + "backend": "tokenizers", + "bos_token": null, + "clean_up_tokenization_spaces": false, + "eos_token": "<|endoftext|>", + "errors": "replace", + "is_local": false, + "local_files_only": false, + "model_max_length": 131072, + "pad_token": "<|endoftext|>", + "padding_side": "right", + "split_special_tokens": false, + "tokenizer_class": "Qwen2Tokenizer", + "unk_token": null +}