Multi-Attribute Constraint Satisfaction via Language Model Rewriting

TL;DR

This work tackles multi-attribute constraint satisfaction for sequential data by introducing MACS, a generalized offline framework that treats a language model as an editor. MACS trains editors on offline edit-pairs using supervised fine-tuning or reward-weighted behavior cloning and guides inference with a reward-based, prioritized editing strategy, enabling real-valued control over multiple attributes. The authors validate MACS on FineCS, a two-task benchmark spanning Text Style Transfer and Protein Design, showing superior constraint satisfaction and the ability to discover novel sequences beyond the training data. The approach offers a flexible, domain-agnostic path toward fine-grained controllability in NLP and bioinformatics, with efficient offline training and broad external-evaluator compatibility.

Abstract

Obeying precise constraints on top of multiple external attributes is a common computational problem underlying seemingly different domains, from controlled text generation to protein engineering. Existing language model (LM) controllability methods for multi-attribute constraint satisfaction often rely on specialized architectures or gradient-based classifiers, limiting their flexibility to work with arbitrary black-box evaluators and pretrained models. Current general-purpose large language models, while capable, cannot achieve fine-grained multi-attribute control over external attributes. Thus, we create Multi-Attribute Constraint Satisfaction (MACS), a generalized method capable of finetuning language models on any sequential domain to satisfy user-specified constraints on multiple external real-value attributes. Our method trains LMs as editors by sampling diverse multi-attribute edit pairs from an initial set of paraphrased outputs. During inference, LM iteratively improves upon its previous solution to satisfy constraints for all attributes by leveraging our designed constraint satisfaction reward. We additionally experiment with reward-weighted behavior cloning to further improve the constraint satisfaction rate of LMs. To evaluate our approach, we present a new Fine-grained Constraint Satisfaction (FineCS) benchmark, featuring two challenging tasks: (1) Text Style Transfer, where the goal is to simultaneously modify the sentiment and complexity of reviews, and (2) Protein Design, focusing on modulating fluorescence and stability of Green Fluorescent Proteins (GFP). Our empirical results show that MACS achieves the highest threshold satisfaction in both FineCS tasks, outperforming strong domain-specific baselines. Our work opens new avenues for generalized and real-value multi-attribute control, with implications for diverse applications spanning NLP and bioinformatics.

Figures (9)

• Figure 1: MACS framework starts with sequential domain datasets (customer reviews or proteins) and a set of real-value attribute evaluators (such as sentiment, complexity regressors, or protein folding models). We then define fine-grained threshold-window boundaries for every attribute and create edit pairs distributed across the multi-attribute landscape. We train the LM editor on top of the edit pairs by leveraging supervised fine-tuning (SFT) or reward-weighted behavior cloning (wBC). Subsequently, LM editors can achieve the desired fine-grained constraints by employing prioritized editing that maintains a priority queue of past edits ordered by their proximity to the target threshold constraints.
• Figure 2: Sentiment and Complexity attribute edit pair distribution via random sampling vs. proposed k-NN sampling. k-NN sampling yields a much more diverse set of edit pairs better suited for simulating editing in all directions.
• Figure 3: Comparing best-of-N vs. reward-prioritized inference constraint satisfaction rate of Sentiment and Complexity attributes. Takeaway: Reward-prioritized inference has better satisfaction rates in harder to reach constraints i.e. edges of the satisfaction matrix.
• Figure 4: Showing 8 attributed paraphrases of a test review for various thresholds generated by + SFT + wBC model with reward-prioritized rewriting.
• Figure 5: Analyzing new mutant discovery of reward-prioritized walk compared with training set distribution. Takeaway: Even with very few training instances in many of the regions, LMs trained with MACS discover many novel candidates. Analysis of the top 15 mutational hotspots reveals that LMs can extrapolate beyond the mutational patterns seen in the training set.