Grammar-State Aware Beam Search for Enhancing Structural Diversity in LLM Generation

Abstract

In many large language model (LLM) applications, generating structurally diverse candidate outputs is crucial for downstream decision-making and reasoning. However, conventional beam search allocates search budget at the token-prefix level, so multiple beam slots may be occupied by lexically different but structurally similar hypotheses. To address these limitations, we propose Grammar-State Aware Beam Search (GSA-Beam). Our method first organizes search at the constraint state level rather than individual prefix level, enabling explicit control over structural diversity. It further incorporates a dynamic state-level beam regulation mechanism that adaptively adjusts beam allocation based on the branching of constraint states, efficiently exploring the feasible structured solution space. Experiments on JSON-schema constrained generation show that GSA-Beam improves explicit state coverage and final structural diversity while preserving schema validity; across several LLM backbones, it reduces reference latency by 22.5% on average compared with standard constrained beam search. Code is available at https://github.com/Paradozile/GSABeam.