Gemma-4-12B-HybridAgent

Gemma-4-12B-HybridAgent is a full-parameter fine-tuned hybrid reasoning agent built upon google/gemma-4-12B (12.5B Total / 11.9B Language Backbone), adapting the Stanford / Together AI HybridAgent (arXiv:2610.18861) framework to the Gemma 4 architecture (Gemma4UnifiedForConditionalGeneration).

By combining Explicit Text Reasoning (<think>...</think>), Continuous Latent Reasoning (<latent>...</latent> via a lightweight residual HybridLink bridge adapter.pt), and Multi-Turn Tool Calling (<python>...</python>, <search>...</search>), the model compresses routine intermediate reasoning into continuous latent vectors while preserving external tool execution and full white-box latent interpretability.

---

⚠️ Important Usage Notice (Read Before Loading!)

Why you MUST NOT use standard AutoModelForCausalLM.from_pretrained(...).generate() directly:

1. adapter.pt & adapter_config.json are HybridLink (29.51M params), NOT a PEFT LoRA adapter! If peft is installed in your environment, calling AutoModelForCausalLM.from_pretrained() directly will mistakenly treat adapter_config.json ({"adapter_type": "ln_res_adapter"}) as a LoRA config and raise a ValueError. Always load via load_hybrid_agent(...).
2. Standard .generate() does not execute latent recurrence: When the model emits <latent> (token_id=262144), the engine must route last-layer hidden states $h_t$ through HybridLink ($\mathcal{R}(h_t)$), scale by Gemma 4's embed_scale ($\sqrt{3840} \approx 61.9677$), and feed the continuous vector back into the decoder until the nearest-token cosine readout hits </latent> (token_id=262145). Use HybridEngine or HybridGenerationEngine below.

---

🛠️ Step 1: Installation

Install the Gemma 4–adapted HybridAgent package:

# Method 1: Clone from GitHub (gemma4 branch)
git clone -b gemma4 https://github.com/milkshake0721/HybridAgent.git
cd HybridAgent
pip install -r requirements.txt
pip install -e .

# Method 2: Or download the bundled hybrid_agent package directly from Hugging Face
python3 -c "from huggingface_hub import snapshot_download; p = snapshot_download('milkshake721/Gemma-4-12B-HybridAgent', allow_patterns=['hybrid_agent/**', 'interpretable_latent_mode/**', 'hybrid_inference/**', 'pyproject.toml', 'requirements.txt']); print('Downloaded package to:', p)"

(Optional) If you want to enable the live web search tool (<search>...</search>), set your Tavily API key:

export TAVILY_API_KEY="tvly-..."

---

🚀 Step 2: Quick Start (4 Ways to Run)

Option A — High-Level Single-Question API (HybridEngine)

import torch
from hybrid_agent.model import load_hybrid_agent
from hybrid_agent.engine import HybridEngine
from hybrid_agent.prompts import SYSTEM_PROMPT

MODEL_ID = "milkshake721/Gemma-4-12B-HybridAgent"

model, tokenizer = load_hybrid_agent(
    MODEL_ID,
    dtype=torch.bfloat16,
    device="cuda",
)

engine = HybridEngine(
    model,
    tokenizer,
    max_new_tokens=1024,
    max_latent_steps=64,
    temperature=0.0,
    top_p=1.0,
)

traj = engine.run(
    question="If 3x + 7 = 22, what is the value of 5x - 2?",
    system_prompt=SYSTEM_PROMPT,
)

print("Final Answer:", traj.final_answer)
print("Total Latent Steps:", traj.latent_steps, "| Latent Spans:", traj.latent_spans)
print("\n--- Rendered Trajectory (with Latent Pseudo-Text Readout) ---")
print(traj.output_rendered)

---

Option B — Full Multi-Turn Lockstep Engine (HybridGenerationEngine with Python & Search Tools)

import torch
from hybrid_agent import (
    HybridGenerationEngine,
    SYSTEM_PROMPT,
    build_tool_executor,
    extract_answer,
    load_hybrid_agent,
)

MODEL_ID = "milkshake721/Gemma-4-12B-HybridAgent"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model, tokenizer = load_hybrid_agent(MODEL_ID, dtype=torch.bfloat16, device=device)
tools = build_tool_executor(require_search=False)

question = "How many trailing zeros does 100! have?"
engine = HybridGenerationEngine(
    model,
    tokenizer,
    device,
    tools,
    SYSTEM_PROMPT,
    question,
    batch_size=1,
    max_response_tokens=1024,
    min_latent_steps=4,
    max_latent_steps=64,
    temperature=0.0,
    top_p=1.0,
    record_latent_readout=True,
)
engine.run(torch.Generator(device=device).manual_seed(0))

res = engine.row_result(engine.rows[0])
print("Answer:", extract_answer(res["output"]))
print("Latent steps:", res["latent_steps"], "| Spans:", res["latent_spans"])
print(res.get("output_rendered", res["output"]))

---

Option C — Interactive White-Box Latent Inspector (interpretable_latent_mode)

Every continuous <latent> thought vector produced by HybridLink can be mapped onto its nearest vocabulary token embedding via cosine similarity, rendering an interactive HTML report that shades each latent step by readability:

python -m interpretable_latent_mode \
    --model milkshake721/Gemma-4-12B-HybridAgent \
    --question "At the speed of sound in air at 20 degrees Celsius, how many seconds does sound take to travel 10 kilometers? Round to the nearest second." \
    --readout_top_k 5 \
    --output latent_report_12b.html

---

Option D — Zero-Clone Standalone Loader (transformers + huggingface_hub)

import torch
import torch.nn as nn
from unittest import mock
from pathlib import Path
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer
import transformers.modeling_utils as modeling_utils

class HybridLink(nn.Module):
    def __init__(self, hidden_size: int):
        super().__init__()
        self.proj1 = nn.Linear(hidden_size, hidden_size)
        self.act = nn.GELU()
        self.proj2 = nn.Linear(hidden_size, hidden_size)
        self.pre_ln = nn.LayerNorm(hidden_size)
        self.post_ln = nn.LayerNorm(hidden_size)

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        return self.post_ln(h + self.proj2(self.act(self.proj1(self.pre_ln(h)))))

repo_dir = Path(snapshot_download("milkshake721/Gemma-4-12B-HybridAgent"))
tokenizer = AutoTokenizer.from_pretrained(str(repo_dir))

with mock.patch.object(
    modeling_utils, "maybe_load_adapters",
    lambda path, download_kwargs, **kw: (None, path, kw),
):
    backbone = AutoModelForCausalLM.from_pretrained(
        str(repo_dir), dtype=torch.bfloat16, device_map="cuda"
    )

hidden_size = backbone.get_input_embeddings().weight.shape[-1]  # 3840
hybridlink = HybridLink(hidden_size).to(device="cuda", dtype=torch.bfloat16)
hybridlink.load_state_dict(torch.load(repo_dir / "adapter.pt", map_location="cuda", weights_only=True))
backbone.eval()
hybridlink.eval()
print("Loaded Gemma-4-12B backbone + HybridLink successfully!")

---

⚙️ Architecture & Training Specifications

Attribute	Value	Notes
Base Model	google/gemma-4-12B	Gemma4UnifiedForConditionalGeneration (48 layers, d_model = 3840)
Language Backbone	11.9B parameters (23.9 GB BF16)	Full-parameter fine-tuned (embed_vision & embed_audio frozen)
Gemma 4 embed_scale	$\sqrt{3840} \approx 61.9677$	Applied automatically in HybridGenerationEngine
HybridLink Adapter (adapter.pt)	29,507,584 parameters (59.0 MB)	2-layer GELU Residual Adapter + LayerNorm (3840 -> 3840)
Expanded Special Tokens	<latent> (262144), </latent> (262145), <think> (262146), </think> (262147)	Registered as atomic special tokens
Training Recipe	3-Stage HybridSFT (format_activation $\rightarrow$ hybridlink_activation $\rightarrow$ self_conditioned)	Trained on togethercomputer/HybridAgent-HybridSFT-v2 (pure BF16)

---

🇹🇼 繁體中文使用說明 (Traditional Chinese Guide)

Gemma-4-12B-HybridAgent 是基於 google/gemma-4-12B 進行全參數微調（Full-Parameter Fine-Tuning）的混合潛空間推理代理模型，將 Stanford / Together AI HybridAgent (arXiv:2610.18861) 架構遷移至 Gemma 4 12B（48 層稠密架構）。

為什麼不能直接用 AutoModelForCausalLM + model.generate()？

1. adapter.pt 與 adapter_config.json 是外掛在模型頭尾之間的 HybridLink 空間橋接器（29.51M 參數），不是 LoRA！ 若直接呼叫 AutoModelForCausalLM.from_pretrained()，環境中的 peft 套件會誤將 adapter_config.json 當成 LoRA 設定檔而報錯。
2. 隱空間思考需要專用狀態機引擎：當模型在文字模式吐出 <latent> (262144) 時，推論引擎必須將最後一層隱藏態 $h_t$ 送入 HybridLink 轉成連續嵌入向量（並乘上 Gemma 4 特有的縮放係數 embed_scale = $\sqrt{3840} \approx 61.9677$），在隱空間高速遞迴推演直到最近鄰餘弦相似度對齊 </latent> (262145) 後再切回文字模式。請使用上方提供的 HybridEngine 或 HybridGenerationEngine。

---

📚 Citation

@misc{zou2026hybridagent,
      title={Hybrid Reasoning Agents}, 
      author={Jiaru Zou and Yunzhe Qi and Sirui Chen and Federico Bianchi and Shang Zhu and Pan Lu and Jingrui He and James Zou},
      year={2026},
      eprint={2610.18861},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2610.18861}, 
}