Guiding Text Generation with Constrained Beam Search in Hugging Face Transformers
Hugging Face has implemented constrained beam search in the transformers library, enabling users to exert precise control over model output by forcing the inclusion of specific words or phrases. This capability is critical for tasks like Neural Machine Translation, where dictionary lookups can dictate mandatory terminology, or when specific formality levels (e.g., formal vs. informal German) must be enforced based on context.
The Challenge of Constrained Generation
Forcing specific subsequences into a text generation task is non-trivial because beam search operates token-by-token. A standard decoding function predicts the next token based on the current sequence but has no inherent mechanism to know if a required token should appear at the current step or at some future step.
Furthermore, managing multiple constraints—such as requiring two different phrases simultaneously or allowing the model to choose one phrase from a list of several possibilities—adds significant complexity to the decoding process.
Constrained Beam Search Capabilities
Constrained beam search allows users to inject prior knowledge into the generation stage rather than filtering outputs after they are produced. The implementation supports two primary types of constraints:
Forced Word Constraints
Using the force_words_ids keyword argument in model.generate(), users can specify exact tokens that must appear in the output. For example, in an English-to-German translation of "How old are you?", forcing the word "Sie" ensures the model produces the formal translation ("Wie alt sind Sie?") instead of the informal one ("Wie alt bist du?").
Disjunctive Constraints
Disjunctive constraints allow the model to satisfy a requirement by including at least one word from a provided list. This is useful when multiple word forms (e.g., "raining", "rained", "rains") are acceptable. In a mixed-constraint scenario, a model can be forced to include one specific word verbatim while choosing one option from a flexible list of related terms.
Technical Implementation: Banks and Round-Robin Selection
Constrained beam search works by injecting desired tokens at every step of generation to move the sequence closer to fulfilling the constraints. However, blindly forcing tokens often leads to nonsensical results (e.g., "The is fast"). To solve this, Hugging Face utilizes a "Banks" system to balance constraint fulfillment with linguistic probability.
The Bank System
Beams are sorted into banks based on their progress toward fulfilling constraints:
- Bank $n$: Contains beams that have made $n$ steps of progress in fulfilling the constraints.
Round-Robin Selection
To prevent the model from only selecting the most constrained (but potentially nonsensical) beams, the algorithm performs a round-robin selection across banks. If num_beams=3, the system selects the most probable output from Bank 2, then Bank 1, then Bank 0. This ensures that high-probability, sensible sequences are preserved even while the model is being pushed toward the constraints.
If a beam deviates from a constraint (e.g., generating "slow" when the constraint was "is fast"), its progress is reset, and it moves back to Bank 0.
Extending Constraints via the Constraint Class
While force_words_ids is the primary user interface, the backend utilizes a Constraint object to track progress and suggest the next tokens. Developers can create custom constraints by subclassing the Constraint abstract interface class.
Potential future extensions mentioned include:
- OrderedConstraints: Ensuring constraints are fulfilled in a specific sequence.
- TemplateConstraints: Forcing the model to follow a specific structural template with gaps to be filled.
Research Foundations
This feature is grounded in several research papers focusing on lexically constrained decoding and guided generation, including work on open vocabulary image captioning and neural machine translation.