Nudging: Inference-time Alignment of LLMs via Guided Decoding

Abstract: LLMs require alignment to effectively and safely follow user instructions. This process necessitates training an aligned version for every base model, resulting in significant computational overhead. In this work, we propose NUDGING, a simple, training-free algorithm that aligns any base model at inference time using a small aligned model. NUDGING is motivated by recent findings that alignment primarily alters the model's behavior on a small subset of stylistic tokens (e.g., discourse markers). We find that base models are significantly more uncertain when generating these tokens. Building on this insight, NUDGING employs a small aligned model to generate nudging tokens to guide the base model's output during decoding when the base model's uncertainty is high, with only a minor additional inference overhead. We evaluate NUDGING across 3 model families on a diverse range of open-instruction tasks. Without any training, nudging a large base model with a 7x-14x smaller aligned model achieves zero-shot performance comparable to, and sometimes surpassing, that of large aligned models. By operating at the token level, NUDGING enables off-the-shelf collaboration between model families. For instance, nudging Gemma-2-27b with Llama-27b-chat outperforms Llama-2-70b-chat on various tasks. Overall, our work offers a modular and cost-efficient solution to LLM alignment. Our code and demo are available at: https://fywalter.github.io/nudging/ .

• The paper presents Nudging, a method that uses token-level interventions from smaller aligned models to guide large language models without additional training.
• It achieves zero-shot performance on 13 tasks by strategically deploying 'nudging tokens' from 7×-14× smaller models to correct uncertainty in larger models.
• The approach fosters cross-family model collaboration, optimizing computational efficiency and ensuring robust performance in reasoning, instruction following, and safety benchmarks.

Nudging: Inference-Time Alignment via Model Collaboration

The paper "Nudging: Inference-time Alignment via Model Collaboration" by Yu Fei, Yasaman Razeghi, and Sameer Singh, presents a novel approach to aligning LLMs at inference time using smaller, aligned models. As LLMs continue to require significant computational resources for alignment across model sizes, this research introduces a method called "Nudging" to address this challenge without additional training.

The Nudging algorithm is motivated by the understanding that alignment often modifies a model's behavior at a superficial level, affecting only a small set of stylistic tokens. The base models exhibit higher uncertainty when generating these tokens. By employing a small aligned model to generate specific "nudging tokens," the algorithm strategically steers the output of larger base models towards desired, aligned behaviors when uncertainty is detected. This research examines nudging's effectiveness across various model families and tasks, demonstrating that nudging can achieve, and sometimes surpass, the zero-shot performance of large aligned models.

Key Findings

1. Performance Enhancement: The research evaluated the nudging technique across three model families—Llama-2, Gemma-2, and OLMo—on 13 tasks, spanning reasoning, general knowledge, instruction following, and safety benchmarks. Without additional training, nudging a base model with a 7×-14× smaller aligned model yielded zero-shot performance comparable to large aligned models, exemplifying significant computational savings.
2. Cross-Family Collaboration: Unlike previous methods, nudging facilitates collaboration between models from different families, enhancing flexibility and demonstrating adaptability for new model families without needing extensive retraining.
3. Token-Level Collaboration: The approach leverages token-level adjustments, using less than 9% of tokens to make a significant 10% improvement on diverse tasks. This precise token collaboration highlights the potential to achieve alignment with minimal intervention.
4. Scaling Analysis: The study reveals marginal benefits from scaling up the nudging model while scaling the base model yields substantial gains, affirming that core model abilities primarily arise during pretraining.
5. Impact on Safety: Nudging proves effective for instruction following and maintaining model safety, achieving aligned-model-level performance across these tasks as validated by GPT-4 evaluations.

Implications and Future Directions

Nudging opens avenues for modular and efficient AI systems, minimizing alignment costs traditionally associated with LLM tuning. The research emphasizes the potential of model collaborations not constrained by family-specific architectures, providing a blueprint for future AI model deployments.

Possible future directions include optimizing nudging rule sets, training specialized smaller nudging models for even greater performance gains, and extending the nudging approach to encompass even larger model ecosystems. By operating at token-level granularity, nudging showcases a simple yet versatile model collaboration framework, essential for next-generation AI design.

Through innovative utilization of smaller aligned models, this work fundamentally contributes to efficient AI model alignment, offering practical and theoretical advancements that promise reduced computational burdens and enhanced AI adaptability.