Air-Decoding: Attribute Distribution Reconstruction for Decoding-Time Controllable Text Generation
Tianqi Zhong, Quan Wang, Jingxuan Han, Yongdong Zhang, Zhendong Mao
Abstract
Controllable text generation (CTG) aims to generate text with desired attributes, and decodingtime-based methods have shown promising performance on this task. However, in this paper, we identify the phenomenon of Attribute Collapse for the first time. It causes the fluency of generated text to rapidly decrease when the control strength exceeds a critical value, rendering the text completely unusable. This limitation hinders the effectiveness of decoding methods in achieving high levels of controllability. To address this problem, we propose a novel lightweight decoding framework named Air-Decoding. Its main idea is reconstructing the attribute distributions to balance the weights between attribute words and nonattribute words to generate more fluent text. Specifically, we train prefixes by prefix-tuning to obtain attribute distributions. Then we design a novel attribute distribution reconstruction method to balance the obtained distributions and use the reconstructed distributions to guide language models for generation, effectively avoiding the issue of Attribute Collapse. Experiments on multiple CTG tasks prove that our method achieves a new state-of-the-art control performance 1 .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 9060a7e0-07cc-42e6-9cb3-d0a95ec01540Cited by top-tier papers6
- Reversing the Forget-Retain Objectives: An Efficient LLM Unlearning Framework from Logit DifferenceJiabao Ji, Yujian Liu, Yang Zhang, Gaowen Liu et al.NeurIPS 2024 · 106 citations
- Controlling Large Language Models Through Concept Activation VectorsHanyu Zhang, Xiting Wang, Chengao Li, Xiang Ao et al.AAAI 2025 · 26 citations
- CoSec: On-the-Fly Security Hardening of Code LLMs via Supervised Co-decodingDong Li, Meng Yan, Yaosheng Zhang, Zhongxin Liu et al.ISSTA 2024 · 10 citations
- Don't Let It Fade: Preserving Edits in Diffusion Language Models via Token Timestep AllocationWoojin Kim, Jaeyoung DoNeurIPS 2025 · 2 citations
- C³TG: Conflict-aware, Composite, and Collaborative Controlled Text GenerationYu Li, Zhe Yang, Yi Huang, Xin Liu et al.AAAI 2026
Builds on14
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- Plug and Play Language Models: A Simple Approach to Controlled Text GenerationSumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung et al.ICLR 2020 · 1,166 citations
- A Contrastive Framework for Neural Text GenerationYixuan Su, Tian Lan, Yan Wang, Dani Yogatama et al.NeurIPS 2022 · 349 citations
- The Power of Scale for Parameter-Efficient Prompt TuningBrian Lester, Rami Al-Rfou, Noah ConstantEMNLP 2021 · 94 citations
- Controlled Text Generation as Continuous Optimization with Multiple ConstraintsSachin Kumar, Eric Malmi, Aliaksei Severyn, Yulia TsvetkovNeurIPS 2021 · 91 citations
Related papers
- Controllable Text Generation via Probability Density Estimation in the Latent SpaceYuxuan Gu, Xiaocheng Feng, Sicheng Ma, Lingyuan Zhang et al.ACL 2023 · 8 citations
- ECO Decoding: Entropy-Based Control for Controllability and Fluency in Controllable Dialogue GenerationSeungmin Shin, Dooyoung Kim, Youngjoong KoEMNLP 2025
- A Distributional Lens for Multi-Aspect Controllable Text GenerationYuxuan Gu, Xiaocheng Feng, Sicheng Ma, Lingyuan Zhang et al.EMNLP 2022 · 16 citations
- DuNST: Dual Noisy Self Training for Semi-Supervised Controllable Text GenerationYuxi Feng, Xiaoyuan Yi, Xiting Wang, Laks V. S. Lakshmanan et al.ACL 2023 · 2 citations
- Controlled LLM Decoding via Discrete Auto-regressive BiasingPatrick Pynadath, Ruqi ZhangICLR 2025
