ABS: Enforcing Constraint Satisfaction On Generated Sequences Via Automata-Guided Beam Search

TL;DR

This paper presents ABS, an automata-guided beam search algorithm that guarantees generated sequences satisfy formal constraints without retraining, improving both compliance and output quality across various tasks.

Contribution

ABS introduces a model-agnostic, inference-time method that enforces constraints via automata-guided beam search, ensuring compliance without retraining models.

Findings

Guarantees constraint satisfaction in sequence generation.

Improves output quality while enforcing constraints.

Effective across multiple tasks including classification and text generation.

Abstract

Sequence generation and prediction form a cornerstone of modern machine learning, with applications spanning natural language processing, program synthesis, and time-series forecasting. These tasks are typically modeled in an autoregressive fashion, where each token is generated conditional on the preceding ones, and beam search is commonly used to balance exploration and fluency during decoding. While deep learning models and Large Language Models (LLMs) excel at capturing statistical patterns in this setting, they remain ill-equipped to guarantee compliance with formal constraints. In this paper, we introduce ABS: a general and model-agnostic inference-time algorithm that guarantees compliance with any constraint that can be compiled into a Deterministic Finite Automaton (DFA), without requiring retraining. ABS leverages the DFA to guide a constrained variant of beam search: at each decoding step, transitions leading to violations are masked, while remaining paths are dynamically re-ranked according to both the model's probabilities and the automaton's acceptance structure. We formally prove that the resulting sequences are guaranteed to satisfy the given constraints, and we empirically demonstrate that ABS also improves output quality. We validate our approach on three distinct tasks: constrained image-stream classification, controlled text generation, and text infilling. In all settings, ABS achieves perfect constraint satisfaction, while outperforming or matching state-of-the-art baselines on standard quality metrics and efficiency.