Source-linked AI summary
Compiling Autoregressive LLM Control into Policy-Exact Parallel Inference for Hierarchical RAG
Zehao Lin, Qianfeng Shen, Xixuan Hao, Shaobo Cui, Chunyu Li, Kai Chen, Bo Tang, Zhenjun Wang, Jin Yan, Zhiyu Li, Feiyu Xiong
TL;DR
Hierarchical RAG controllers face a choice between failure-prone free-form JSON, history-blind unary selection, and serial grammar-constrained decoding. This paper compiles bounded-horizon action prefixes into parallel probes that reproduce the serial policy in exact arithmetic. Across splits, compiled control matches serial action sequences on most questions while reducing control latency.
Problem
Hierarchical RAG uses input-dependent section choices, but free-form JSON risks protocol errors, unary selection loses history dependence, and grammar-constrained decoding requires serial forwards.
Source: §1 · p. 1Method
The method compiles a bounded-horizon action-prefix tree into a root prefill and tree-causal probe that evaluates the serial controller’s prefix-conditioned distributions.
Source: §1 · p. 1Results
Compiled control matches serial action sequences on 84–89% of questions per split (87% at 3B) and cuts mean control latency by 35–41% (47% at 3B).
Source: §4.1 · p. 3Takeaways & Limitations
The construction applies to finite-action constrained LLM controllers, with hierarchical RAG serving as the case study.
Source: §1 · p. 1Different floating-point reduction orders can flip near-tie slots, although every observed divergence fell within the stated deviation window.
Source: §3.3 · p. 3
Abstract
from arXiv · showhide
Large language models increasingly act as controllers in structured decision pipelines, yet their decisions are still executed by autoregressive text decoding. Hierarchical RAG is one such setting: at each frontier of a document tree, an LLM commits to $b$ of $k$ sections under an input-dependent action grammar, executed as either failure-prone free-form JSON or reliable but strictly serial grammar-constrained generation. We show that this serialization is unnecessary for bounded-horizon control: the reachable future forms a shallow action-prefix tree, which we compile --- one root prefill plus one tree-causal probe evaluate every prefix-conditioned action distribution with serial decoding's tokens, positions, and causal ancestry. A prefix-state lemma makes compiled control policy-exact in exact arithmetic: the same policy with a two-forward critical path and zero output tokens. One further probe pass certifies the exact order-marginalized MAP set. Across four document-disjoint splits and two scales on the primary MetaX device, compiled control matches the serial policy on $84$--$89\%$ of questions --- every instrumented divergence a floating-point near-tie --- and cuts control latency by $35$--$47\%$ versus token-wise serial decoding and $6$--$7\%$ versus an idealized chunked engine; an H800 repetition preserves the ordering.
1. INTRODUCTION
Hierarchical RAG makes evidence gathering a sequential control problem, but free-form JSON, unary selection, and grammar-constrained decoding each impose a different limitation. The paper compiles bounded-horizon action prefixes to preserve the constrained policy while enabling parallel evaluation and exact set inference.
1. INTRODUCTION
- Hierarchical RAG selects sections from input-dependent candidate sets, while JSON, unary selection, and grammar-constrained decoding respectively risk protocol errors, lose selection history, or require serial forwards.The companion interface study reports 3–65% of questions failing on syntax, identifiers, or cardinality in free-form JSON.Source: §1 · p. 1
- A bounded horizon limits reachable states to action prefixes, so one root prefill and one tree-causal probe evaluate every prefix-conditioned distribution with the serial controller’s tokens, positions, and causal ancestry.The compiled critical path has two dependent forwards, independent of horizon, with zero generated tokens and no protocol surface.Source: §1 · p. 1
- The probe also supports order-marginalized set inference with a branch-and-bound MAP certificate and bounds the unary shortcut’s expected-utility loss via AR-vs-unary KL divergence.The construction applies to finite-action constrained controllers; hierarchical RAG is the case study.Source: §1 · p. 1
2. RELATED WORK
Prior work frames hierarchical retrieval as document-structure-guided evidence gathering and establishes constrained decoding and tree-parallel execution as related approaches. This paper’s distinction is to exhaustively evaluate a frozen controller’s reachable conditionals rather than approximate joint structure or accelerate speculative decoding.
2. RELATED WORK
- Tree- and memory-structured retrieval steer an LLM through document structure, while evidence quality is sensitive to structure and OCR noise.The passage cites tree- and memory-structured retrieval [2] [4] and evidence quality sensitivity.Source: §2 · pp. 1–2
- Grammar-constrained decoding preserves validity but keeps a serial path, whereas first-token logit selection avoids decoding but abandons sequential policy and inherits label-position sensitivities.The related-work passage cites constrained decoding [11] [12] and logit selection, with label-position sensitivities.Source: §2 · p. 2
- Unlike speculative decoding, compilation evaluates the exhaustive reachable state space of a bounded-horizon grammar, so there is nothing to reject and the equivalence is provable rather than distributional.Figure 1 depicts serial decoding’s dependent forwards alongside a root prefill and tree-causal probe; it reports a fresh-split latency reduction and zero output tokens and protocol failures.Source: §2 · p. 2, Figure 1 · p. 2
3. COMPILED AUTOREGRESSIVE CONTROL
The paper replaces serial constrained decoding with a two-step compilation of the controller’s reachable action prefixes. In exact arithmetic, this preserves the serial policy while supporting certified set inference and an interaction-mass diagnostic.
3.1. The serial constrained controller
- The serial controller emits a greedy sequence of distinct labels under a grammar that masks labels already selected, with each emitted token requiring a dependent forward pass.For three labels, the critical path is approximately nine forwards; even chunking forced spans still requires one forward per label slot.Source: §3.1 · p. 2
3.2. Compilation to an action-prefix tree
- Compilation evaluates reachable action-prefix states with a root prefill and one tree-causal probe, reducing the critical path to two forwards.The probe uses serial position indices and each prefix’s causal ancestors; its parallel width grows with the number of prefix continuations.Source: §3.2 · pp. 2–3
3.3. Policy exactness
- With matching token and position ids, ancestor-visible attention, parameters, masks, candidate order, and tie-breaking, compiled and serial greedy control select the same sequence in exact arithmetic.Hardware’s different floating-point reduction orders can produce data-dependent logit deviations and flip near-tie slots.Source: §3.3 · p. 3
3.4. Certified set inference and interaction mass
- The probe supports exact order-marginalized MAP set inference: expanding ordered pairs by mass can certify the best set without truncation.The interaction-mass diagnostic bounds unary-shortcut expected-utility loss by Iseq/2; in split B, only 14% of rounds had Iseq below 0.5.Source: §3.4 · p. 3
4. EXPERIMENTS
Across the experiments, compiled control closely tracks the serial controller while reducing control latency, with measured gains depending on precision, workload, and probe width. Ablations show that later history-conditioned slots matter, and synthetic sweeps expose a width limit beyond the deployed horizon.
4.1. Compilation reproduces the serial policy
- Compiled control matches serial action sequences on 84–89% of questions per split (87% at 3B), with zero output tokens and 35–41% lower mean control latency (47% at 3B).On D and E, it also beats the chunked baseline: serial/chunked/compiled latency is 0.58/0.37/0.35 s on D and 0.64/0.40/0.37 s on E.Source: §4.1 · pp. 3–4
- All 39 instrumented diverging rounds occur at decision margins ≤2.5, while H800 repetitions preserve latency ordering but show 65–77% final-selection agreement.A full fp32 replay reduces split-A deviations by over 20×; serial decoding itself agrees across accelerators on only 67–73% of questions.Source: §4.1 · p. 4
4.2. Quality tracks the source policy
- The AR policy’s advantage over unary marginals varies by split: +5.9/+4.4 points at 2K on B/C, but −1.1 on D; false-nesting tests show 0 protocol errors for compiled control.On the corrupted hierarchy, compiled control scores 44.4 AllEvidence@1K, while free generation fails on 7.8–14.4% of questions.Source: §4.2 · p. 4
4.3. Ablations
- On split A, first-slot-only selection scores 42.2 versus 51.1 AllEvidence@1K for the full compiled chain, indicating that later history-conditioned slots carry information.A sequential-label serialization also changes the policy, scoring 44.4 versus 50.0 AllEvidence@1K.Source: §4.3 · p. 4
4.4. Work–depth scaling
- At k=20, compiled selection is 1.3–2.8× faster than serial for n≤3, but the n=4 probe never catches the chunked engine and loses at short prompts.At n=4, the probe grows to 362 rows; compilation uses 13× the serial work at n=3 and about 270× at n=4, trading parallel work for a shorter critical path.Source: §4.4 · p. 4, Figure 2 · p. 4
5. CONCLUSION
The conclusion presents compilation as an alternative to decoding bounded-horizon constrained controllers. It identifies probe width and bf16 near-tie divergences as limits, while describing broader applicability to finite-action agent loops.
5. CONCLUSION
- A tree-causal probe evaluates a bounded-horizon controller’s shallow action-prefix tree, reproducing the serial policy in exact arithmetic with a depth-2 critical path and zero output tokens.The conclusion also describes certified set inference and an interaction-mass bound as consequences of the construction.Source: §5 · p. 4
- Probe width grows factorially beyond the greedy horizon, while fp32 narrows bf16 near-tie divergences to 89/90.The probe is an ordinary attention call, and the paper says the guarantees transfer to finite-action agent loops.Source: §5 · p. 4
- The paper says the compiler and benchmark harness will be released, and that controller semantics can be decoupled from the serial critical path.Source: §5 · p. 4
6. COMPLIANCE WITH ETHICAL STANDARDS
The study reports using publicly available QASPER papers and no human subjects, animal experiments, or personally identifiable data.
6. COMPLIANCE WITH ETHICAL STANDARDS
- The study uses publicly available QASPER papers and no human subjects, animal experiments, or personally identifiable data; no ethical approval was required.Source: §6 · p. 5