From Clicks to Preference: A Multi-stage Alignment Framework for Generative Query Suggestion in Conversational System

TL;DR

This paper presents a multi-stage framework for generative query suggestion in conversational systems, aligning model outputs with user preferences through prompt engineering, fine-tuning, probabilistic preference modeling, and reinforcement learning, resulting in improved user engagement.

Contribution

The paper introduces a novel multi-stage alignment framework incorporating Gaussian Reward Models and out-of-distribution regularization for better preference modeling and policy alignment.

Findings

Significant improvement over baselines in automatic and human evaluations.

34% relative increase in user engagement in live A/B tests.

Effective modeling of user preferences as probability distributions.

Abstract

Generative query suggestion using large language models offers a powerful way to enhance conversational systems, but aligning outputs with nuanced user preferences remains a critical challenge. To address this, we introduce a multi-stage framework designed for progressive alignment between the generation policy and user intent. Our pipeline begins with prompt engineering as a cold-start strategy, followed by the Supervised Fine-Tuning stage, in which we introduce a distillation method on click logs to create a robust foundational model. To better model user preferences while capturing their inherent uncertainty, we develop a Gaussian Reward Model (GaRM) that represents user preferences as probability distributions rather than point estimates. Finally, we employ reinforcement learning to align the generation policy with these preferences, guided by a composite reward function that integrates GaRM with auxiliary heuristics to mitigate reward hacking. To maintain training stability, this process is enhanced by a novel out-of-distribution regularization method and a two-stage reward fusion technique. Extensive experiments demonstrate that our framework significantly outperforms baselines on both automatic and human evaluations and yields a 34\% relative increase in user engagement as measured by click-through rate in live A/B tests.