uuugi's picture
Add GitHub repo badges, clone options, and Paper PDF links
6afdcc7 verified
|
Raw
History Blame Contribute Delete
10.7 kB
---
license: mit
tags:
- constrained-decoding
- reachability
- logit-processor
- structured-generation
- grammar-masking
- dfa
- fsm
- pytorch
- transformers
pipeline_tag: text-generation
---
# Goal-Conditioned Reachability Logit Masker (GCLM)
[![GitHub Repo](https://img.shields.io/badge/GitHub-Repository-181717.svg?logo=github)](https://github.com/uuuugi/Goal-Conditioned-Reachability-Logit-Masker)
[![Paper PDF](https://img.shields.io/badge/Paper-PDF-b31b1b.svg?logo=adobeacrobatreader)](paper.pdf)
[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Repository-yellow)](https://huggingface.co/uuugi/gclm-constrained-decoding)
[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
[![PyTorch 2.0+](https://img.shields.io/badge/PyTorch-2.0+-ee4c2c.svg)](https://pytorch.org/)
[![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)
An ultra-fast, strictly **O(1)** runtime **Goal-Conditioned Reachability Logit Masking Engine** for Large Language Models.
GCLM mathematically guarantees that an LLM will strictly reach designated goal/accepting states within a fixed token budget (`T_max`), **fundamentally preventing dead-end traps and truncated syntax failures**.
📄 **Paper**: [**Read / Download `paper.pdf`**](paper.pdf)  |  💻 **GitHub**: [**uuuugi/Goal-Conditioned-Reachability-Logit-Masker**](https://github.com/uuuugi/Goal-Conditioned-Reachability-Logit-Masker)
---
## 💡 Key Differences: GCLM vs. Forward DFA Maskers (Outlines / SGLang)
```
[Traditional Forward DFA (Outlines / SGLang)]
Start (A) ─── Token X ───▶ [Valid Branch D] ─── Token Y ───▶ [Dead-End / Truncated Trap ❌]
(Only checks if transition exists from current state)
[GCLM: Time-Bounded Backward Reachability (Ours)]
Start (A) ─── Token X (Masked to -inf ⛔)
└── Token B ───▶ State C ───▶ Goal / Closing '}' ✅
(Preemptively prunes any branch that cannot reach Goal in <= T_rem steps)
```
| Feature | Standard Forward DFA (Outlines / SGLang) | **GCLM (Ours)** |
| :--- | :--- | :--- |
| **Masking Basis** | Current state validity (`s_curr -> s'`) | **Time-bounded backward reachability** (`s_curr -> s' ->* S_goal` in `≤ T_rem - 1` steps) |
| **Dead-End Traps** | ❌ May enter valid forward branches that lead to dead-ends | ✅ **Preemptively masked** before entering trap |
| **Token Budget Exceeded** | ❌ Outputs truncated/broken syntax when budget ends | ✅ **Forces early syntax closure** before budget exhaustion |
| **Per-Token Overhead** | O(1) table lookup | **Strict O(1) vectorized PyTorch lookup (< 0.1ms)** |
| **Complexity Scaling** | Scales with active state transitions | **Zero runtime dependence on state count (S)** |
---
## 📐 Mathematical Formulation
### 1. Offline Backward BFS Table Builder
Given an FSM `(S, Σ, δ, s_0, S_goal)` and maximum token budget `T_max`, we precompute a reachability tensor `R` of shape `(T_max + 1, |S|)` via vectorized backward BFS:
```python
# Base Step (t = 0):
R[0, s] = 1 if (s in S_goal) else 0
# Vectorized Backward BFS (for t = 1 ... T_max):
R[t, s] = R[t-1, s] OR (∃ v ∈ V such that δ(s, v) >= 0 and R[t-1, δ(s, v)] == 1)
```
### 2. Strict O(1) Runtime Logits Masking
At decoding step `k` with remaining token budget `T_rem = T_max - k`:
```python
# Step 1: Vectorized check for valid transitions within remaining budget
ValidTokens(v) = (δ(s_curr, v) >= 0) AND R[min(T_rem - 1, T_max), clamp(δ(s_curr, v), 0)]
# Step 2: In-place O(1) logit masking
Logits[v] = Logits[v] if ValidTokens(v) == 1 else -inf
```
---
## 📁 Repository Structure
```
gclm_project/
├── core/
│ ├── __init__.py
│ ├── fsm_builder.py # Transitions tensor & vectorized backward BFS reachability table
│ ├── logit_processor.py # Hugging Face LogitsProcessor compatible O(1) in-place masker
│ └── compiler.py # Tokenizer-aware grammar/pattern compiler
├── benchmarks/
│ ├── synthetic_deadend.py # Experiment 1: Dead-end trap avoidance benchmark
│ ├── json_budget_bench.py # Experiment 2: Real-world strict budget JSON benchmark
│ ├── tool_calling_bench.py # Experiment 3: Multi-step agent action budget benchmark
│ ├── scaling_bench.py # Experiment 4: Complexity scaling (|S|=10~10,000) & plot generator
│ ├── real_model_bench.py # Experiment 5: Real lightweight LLM (Qwen2.5) E2E benchmark
│ └── latency_bench.py # Per-token runtime overhead benchmark
├── examples/
│ └── run_generation.py # Live interactive generation demo with Transformers
├── tests/
│ ├── test_fsm_builder.py # Unit tests for BFS reachability & multi-goal
│ └── test_logit_processor.py # Unit tests for batch masking & state progression
├── paper_figure_scaling.png # Publication-ready 300-DPI scaling figure
├── requirements.txt
└── README.md
```
---
## 📊 Comprehensive Experimental Results
### 1. Real Lightweight LLM End-to-End Benchmark (`Qwen2.5-0.5B`)
> Tested on real model weights generating JSON responses under strict token limits.
| Token Budget (T_max) | Vanilla Sampling | Forward DFA (Outlines Style) | **GCLM (Ours)** | Latency / Sample (GCLM) |
| :--- | :---: | :---: | :---: | :---: |
| **T_max = 6 tokens** | 0.0% | 30.0% | **100.0%** | **615.90 ms** (Fastest, early closure) |
| **T_max = 10 tokens** | 0.0% | 70.0% | **100.0%** | **1,086.02 ms** |
| **T_max = 16 tokens** | 0.0% | 85.0% | **100.0%** | **992.39 ms** |
---
### 2. Strict Budget JSON Schema Parsing Benchmark
> Complex nested JSON schema tested across 500 trials per budget.
| Budget (T_max) | Vanilla | Forward DFA (Outlines Style) | **GCLM (Ours)** | Key Insight |
| :--- | :---: | :---: | :---: | :--- |
| **T_max = 4** | 2.4% | 55.4% | **100.0%** | **Forces safe `{}` closure when fields cannot finish** |
| **T_max = 6** | 2.4% | 45.6% | **100.0%** | Prunes deep nested object paths |
| **T_max = 8** | 2.2% | 65.2% | **100.0%** | Eliminates dangling commas |
| **T_max = 16** | 1.4% | 91.8% | **100.0%** | Complete 100% parse rate across all budgets |
---
### 3. Multi-Step Agent Tool-Calling & Action Budget Benchmark
> ReAct-style multi-tool workflow evaluating goal completion within action limits.
| Action Budget | Vanilla | Forward DFA | **GCLM (Ours)** | Key Finding |
| :--- | :---: | :---: | :---: | :--- |
| **3 Actions** | 0.00% | 16.80% | **100.00%** | Dynamically forces 3-step shortest path |
| **4 Actions** | 0.00% | 33.20% | **100.00%** | Prunes unfinishable deep search subtrees |
| **8 Actions** | 0.60% | 65.20% | **100.00%** | **Completely avoids infinite retry trap loops** |
---
### 4. FSM Complexity & Strict O(1) Runtime Scaling
> Scaling state count `|S|` from 10 to 10,000 (1,000x increase). Plot saved as `paper_figure_scaling.png`.
| Vocabulary Size (V) | State Count (S) | Offline BFS Time | Memory Footprint | Online Latency per Token |
| :--- | :---: | :---: | :---: | :---: |
| **V = 32,000 (LLaMA)** | S = 10 | 29.55 ms | 2.44 MB | **388.72 µs** |
| V = 32,000 | S = 100 | 240.10 ms | 24.42 MB | **335.10 µs** |
| V = 32,000 | S = 1,000 | 2,111.82 ms | 244.19 MB | **340.84 µs** |
| V = 32,000 | **S = 10,000** | 25,790.14 ms | 2.44 GB | **356.29 µs** (O(1) verified) |
| **V = 151,643 (Qwen2.5)** | S = 10 | 159.29 ms | 11.57 MB | **601.92 µs** |
| V = 151,643 | **S = 10,000** | 147,702.79 ms | 11.56 GB | **666.22 µs** (O(1) verified) |
---
## 🚀 Quick Start
### 1. Installation
**From Hugging Face:**
```bash
git clone https://huggingface.co/uuugi/gclm-constrained-decoding
cd gclm-constrained-decoding
pip install -r requirements.txt
```
**From GitHub:**
```bash
git clone https://github.com/uuuugi/Goal-Conditioned-Reachability-Logit-Masker.git
cd Goal-Conditioned-Reachability-Logit-Masker
pip install -r requirements.txt
```
### 2. Basic Usage with Hugging Face Transformers
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessorList
from core.fsm_builder import ReachabilityFSM
from core.logit_processor import GoalReachabilityLogitsProcessor
# 1. Load model and tokenizer
model_id = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
vocab_size = model.config.vocab_size
max_budget = 15
# 2. Define FSM & Goal state
fsm = ReachabilityFSM(num_states=5, vocab_size=vocab_size)
fsm.add_transition(from_state=0, token_id=101, to_state=1)
fsm.add_transition(from_state=1, token_id=102, to_state=2)
fsm.set_goal_states([2])
# 3. Precompute reachability table (one-time offline step)
fsm.build_reachability(max_steps=max_budget)
# 4. Attach GCLM to Hugging Face LogitsProcessorList
gclm_processor = GoalReachabilityLogitsProcessor(fsm=fsm, max_budget=max_budget)
logits_processors = LogitsProcessorList([gclm_processor])
# 5. Generate with guaranteed reachability
inputs = tokenizer("Your prompt here", return_tensors="pt")
outputs = model.generate(
**inputs,
max_new_tokens=max_budget,
logits_processor=logits_processors
)
print(tokenizer.decode(outputs[0]))
```
---
## 🧪 Reproducing Experiments
```bash
# Run Unit Tests
python -m pytest tests/ -v
# Run Experiment 1: Synthetic Dead-End Benchmark
python -m benchmarks.synthetic_deadend
# Run Experiment 2: Strict Budget JSON Benchmark
python -m benchmarks.json_budget_bench
# Run Experiment 3: Agent Tool-Calling Benchmark
python -m benchmarks.tool_calling_bench
# Run Experiment 4: Scaling Benchmark & Generate Paper Plots
python -m benchmarks.scaling_bench
# Run Experiment 5: Real Lightweight LLM Benchmark (Qwen2.5)
python -m benchmarks.real_model_bench --model Qwen/Qwen2.5-0.5B
```
---
## 📑 Paper & Citation
📄 **Paper PDF**: [**Download `paper.pdf`**](paper.pdf)
💻 **GitHub Repository**: [**uuuugi/Goal-Conditioned-Reachability-Logit-Masker**](https://github.com/uuuugi/Goal-Conditioned-Reachability-Logit-Masker)
🤗 **Hugging Face Model**: [**uuugi/gclm-constrained-decoding**](https://huggingface.co/uuugi/gclm-constrained-decoding)
```bibtex
@article{an2026gclm,
title={Goal-Conditioned Reachability Logit Masker: Guaranteed Goal Satisfaction for Constrained LLM Generation in O(1) Time},
author={An, ByeongUk},
journal={arXiv preprint},
year={2026}
}
```
**Author**: ByeongUk An
**Email**: `hhjjkk7186@gmail.com`
**ORCID**: [`0009-0007-5612-5602`](https://orcid.org/0009-0007-5612-5602)
---
## 📄 License
MIT License