Qwen3-4B-Instruct-2507-zip-rc/README.md

---
library_name: transformers
license: apache-2.0
base_model: Qwen/Qwen3-4B-Instruct-2507
pipeline_tag: text-generation
tags:
- zip-rc
- adaptive-compute
- introspection
- reasoning
---

# Qwen3-4B-Instruct-2507-ZIP-RC

This model is a modified version of **[Qwen/Qwen3-4B-Instruct-2507](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507)**, trained to support **Zero-Overhead Introspection (ZIP-RC)** for adaptive test-time compute.

It was created as part of a **Paper Replication** experiment for:
**"Zero-Overhead Introspection for Adaptive Test-Time Compute"** (Manvi et al., 2025).

| **Links** | **Description** |
| :--- | :--- |
| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dataopsnick/paper-replication/blob/main/zip-rc_replication/dataopsnick_Qwen3_4B_Instruct_2507_zip_rc_QUICKSTART.ipynb) | **Quickstart Notebook:** Run adaptive inference immediately. |
| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dataopsnick/paper-replication/blob/main/zip-rc_replication/paper_replication_arxiv_org_abs_2512_01457.ipynb) | **Replication Notebook:** Full experiments and reproduction. |

## Model Description

This model retains the full reasoning capabilities of the base `Qwen3-4B-Instruct` model but features a **fine-tuned LM Head**. The head has been trained to repurpose unused logit space to predict a **joint distribution of Expected Reward (Correctness) and Remaining Generation Length** at every token step.

This allows the model to "introspect" during generation with **zero computational overhead**, enabling:
*   **Adaptive Sampling:** Dynamically pruning low-quality trajectories.
*   **Budget Management:** Balancing compute cost vs. accuracy.
*   **Self-Correction:** Detecting when a reasoning path is failing before it finishes.

## Usage

### 1. Quick Start: Adaptive Inference
The easiest way to use the model is via the `ziprc` helper library, which handles the Meta-MDP logic (branching, pruning, and swapping).

```python
import torch
import sys
import os
from huggingface_hub import hf_hub_download

# 1. Download the helper script dynamically
script_path = hf_hub_download(repo_id="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc", filename="ziprc.py")
sys.path.append(os.path.dirname(script_path))

# 2. Import the downloaded module
import ziprc 

# 3. Run Inference
model = ziprc.ZIPRCModel(ziprc.ZIPRCConfig())
sampler = ziprc.ZIPRCSampler(model)

prompt = "Solve the following logic puzzle: Five adults check into a hotel with three dogs. How many shoes are they all wearing?"
trajectories = sampler.generate(prompt, initial_samples=2)

best = sampler.select_best_trajectory(trajectories)
print(f"Confidence: {best['final_score']:.2%}")
```

### 2. Advanced Usage: Streaming & Configuration
This example shows how to configure the pruning aggressiveness (`alpha`) and cost penalty (`beta`), and how to stream the result to see the introspection in action.

```python
import sys
import os
#import tqdm
from huggingface_hub import hf_hub_download
from tqdm import tqdm

# 1. Download the helper script dynamically from the repo
script_path = hf_hub_download(repo_id="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc", filename="ziprc.py")
sys.path.append(os.path.dirname(script_path))

# 2. Import the module
from ziprc import ZIPRCModel, ZIPRCConfig, ZIPRCSampler

# 3. Configure and Load Model
# Note: The model weights are downloaded automatically here
cfg = ZIPRCConfig(
    model_name="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc",
    alpha=0.1,          # Threshold for pruning
    beta=0.05,          # Cost penalty
    smoothing_window=3  # For stable predictions
)

model = ZIPRCModel(cfg)
sampler = ZIPRCSampler(model)

# 4. Generate with Introspection
prompt = "Solve the following logic puzzle: Five adults check into a hotel with three dogs. How many shoes are they all wearing?"

# generate_stream produces trajectories with introspection data
trajectories = sampler.generate_stream(prompt, initial_samples=2)

# Select the best answer based on the introspection score
best = sampler.select_best_trajectory(trajectories)

print(f"Confidence: {best['final_score']:.2%}")
print(f"Answer: {model.tokenizer.decode(best['ids'][0], skip_special_tokens=True)}")
```

### 3. Low-Level: Reading the Logits
You can manually decode the introspection signal (Reward and Cost) from the reserved tokens in the logits without using the sampler.

```python
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "dataopsnick/Qwen3-4B-Instruct-2507-zip-rc"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")

# Configuration used during training
reward_bins = 8
length_bins = 7
total_zip_tokens = 56
zip_start_offset = 56
# ZIP tokens are located at the very end of the vocabulary
zip_start_id = model.config.vocab_size - zip_start_offset

def get_introspection_probs(logits):
    """
    Extracts the joint distribution P(Reward, Length) from the logits.
    """
    # Slice the reserved ZIP logits
    zip_logits = logits[:, zip_start_id : zip_start_id + total_zip_tokens]
    
    # Softmax over the flat ZIP tokens to get valid probabilities
    probs = F.softmax(zip_logits, dim=-1)
    
    # Reshape to [Batch, Reward_Bins, Length_Bins]
    return probs.view(-1, reward_bins, length_bins)

# Example Inference Step
prompt = "Solve the following logic puzzle: Five adults check into a hotel with three dogs. How many shoes are they all wearing?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model(inputs.input_ids)
    next_token_logits = outputs.logits[:, -1, :]
    
    # Get Introspection Signal (Zero Overhead)
    joint_dist = get_introspection_probs(next_token_logits)
    
    # 1. Marginalize over length to get P(Reward) distribution
    p_reward = joint_dist.sum(dim=2)  # Shape: [Batch, Reward_Bins]
    
    # 2. Calculate Expected Reward (Confidence)
    # The reward bins are linearly spaced [0, 1]. We use bin centers for the weighted sum.
    # centers = 0.0625, 0.1875, ..., 0.9375
    reward_grid = torch.linspace(0.0625, 0.9375, reward_bins).to(model.device)
    
    # E[R] = sum(P(r) * r)
    expected_reward = (p_reward * reward_grid).sum(dim=1).item()
    
    print(f"Model Confidence: {expected_reward:.2%}")
```

### 4. OpenAI-Compatible Streaming (Async)
This method exposes the introspection data (`zip_rc` field) alongside standard text generation chunks, suitable for integration with frontends.

```python
import asyncio
import nest_asyncio
from ziprc import ZIPRCModel, ZIPRCConfig, ZIPRCSampler

# 1. Setup (Run once)
# This patch is required for running async loops in Colab/Jupyter
nest_asyncio.apply()

# Load Model
cfg = ZIPRCConfig(model_name="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc")
model = ZIPRCModel(cfg)
sampler = ZIPRCSampler(model)

async def consume_inference_stream():
    prompt = "Solve the following logic puzzle: Five adults check into a hotel with three dogs. How many shoes are they all wearing?"
    
    print(f"User: {prompt}\n" + "-"*60)
    print("Assistant (Streaming with Introspection):")
    
    # 2. Get the OpenAI-compatible stream
    # Returns an async generator yielding chunk objects
    stream = sampler.openai(prompt, max_tokens=256)

    final_clean_answer = ""

    async for chunk in stream:
        # --- Channel A: Standard Text (Compatible with standard UIs) ---
        # Use .get() to handle the final chunk where delta is empty
        # Use .get() to safely handle the final chunk where delta is empty
        delta = chunk.choices[0].delta
        content = delta.get("content", "")
        
        if content:
            print(content, end="", flush=True)
                    
        # --- Channel B: Zero-Overhead Introspection (The "Pareto" Gain) ---
        # We access the side-channel data to see what the model is thinking
        # without running separate reward model inference.
        if hasattr(chunk, 'zip_rc'):
            info = chunk.zip_rc
            
            # If the model performs a meta-action (Branching/Pruning), log it
            # Filter out 'finished' to avoid accessing missing utility/score fields
            if info.action not in ['keep', 'finished']:
                print(f"\n[⚙️ META-ACTION: {info.action} | Utility: {info.utility:.4f}] ", end="")

            # Check for the Final Answer
            if info.get('action') == 'finished' and 'final_text' in info:
                final_clean_answer = info['final_text']
            
            # Optional: Peek at the "Confidence" (Expected Correctness) in real-time
            # if info.step % 10 == 0:
            #     print(f" (Conf: {info.lhs_score:.1%}) ", end="")

    print("\n" + "-" * 40)
    print("🏆 FINAL BEST ANSWER (Clean):")
    print("-" * 40)
    print(final_clean_answer)

# 3. Execution
loop = asyncio.get_event_loop()
loop.run_until_complete(consume_inference_stream())
```

### 5. Local Server Deployment
You can deploy an OpenAI-compatible API server that streams both text and introspection data.

```python
import sys
import os
import asyncio
import uvicorn
from huggingface_hub import hf_hub_download

# 1. Download server.py
script_path = hf_hub_download(repo_id="dataopsnick/Qwen3-4B-Instruct-2507-zip-rc", filename="server.py")
sys.path.append(os.path.dirname(script_path))

# 2. Import the app
# NOTE: This will load the model weights again if they aren't cached.
# If you are low on VRAM, restart your runtime before running this cell.
from server import app

# 3. Run the Server (Colab/Jupyter Compatible)
HOST = "0.0.0.0"
PORT = 8000
config = uvicorn.Config(app, host=HOST, port=PORT)
server = uvicorn.Server(config)

try:
    # Check if we are in an existing loop (Colab)
    loop = asyncio.get_running_loop()
    print(f"🚀 Server running in background on http://{HOST}:{PORT}")
    loop.create_task(server.serve())
except RuntimeError:
    # Standard script execution
    asyncio.run(server.serve())
```

## Citation

```bibtex
@article{manvi2025ziprc,
  title={Zero-Overhead Introspection for Adaptive Test-Time Compute},
  author={Manvi, Rohin and Hong, Joey and Seyde, Tim and Labonne, Maxime and Lechner, Mathias and Levine, Sergey},
  journal={arXiv preprint arXiv:2512.01457},
  year={2025}
}
```