Stream Aligner: Efficient Sentence-Level Alignment via Distribution Induction

TL;DR

The paper tackles the challenge of aligning large language models efficiently across tasks by proposing Stream Aligner, a sentence-level, distribution-inducing correction mechanism that operates during generation. By training a compact model to learn residuals between preferred and original suffix sentences and iteratively feeding corrections back into the generation prefix, Stream Aligner reduces reliance on large auxiliary models while enhancing reasoning and safety. Empirical results show substantial gains in helpfulness and harmlessness for QA and noticeable improvements in math tasks, with performance improving as more correction rounds are applied. The approach achieves a favorable balance between effectiveness and inference-time efficiency, enabling scalable, low-latency alignment across tasks. Overall, Stream Aligner demonstrates that sentence-level, distribution-driven corrections can elicit latent capabilities of the upstream model more effectively than traditional alignment methods.

Abstract

The rapid advancement of large language models (LLMs) has led to significant improvements in their capabilities, but also to increased concerns about their alignment with human values and intentions. Current alignment strategies, including adaptive training and inference-time methods, have demonstrated potential in this area. However, these approaches still struggle to balance deployment complexity and capability across various tasks and difficulties. In this work, we introduce the Streaming Distribution Induce Aligner (Stream Aligner), a novel alignment paradigm that combines efficiency with enhanced performance in various tasks throughout the generation process. Stream Aligner achieves dynamic sentence-level correction by using a small model to learn the preferences of the suffix sentence, iteratively correcting the suffix sentence output by the upstream model, and then using the corrected sentence to replace the suffix sentence in subsequent generations. Compared to Aligner, our experiments demonstrate that Stream Aligner reduces reliance on the capabilities of additional models, enhances the reasoning abilities of LLMs, and decreases latency during user interaction. Specifically, Stream Aligner-2B model has achieved an improvement of 76.1% in helpfulness, 36.0% in harmlessness on the tested Llama2-70B-chat model, and Stream Aligner-8B has achieved an improvement of 3.5% on the math ability of the tested Llama3-70B-Instruct model.

Figures (9)

• Figure 1: Operational Dynamics of the Stream Aligner Generation Pipeline.Stream Aligner serves as an plug-and-play module in the generation pipeline. It corrects the sentence generated by the upstream model and then feeds the corrected suffix back to the upstream model for further generation until the end of responses. This pipeline ensures every sentence in the output is aligned with the Stream Aligner model, thereby aligned with the human preference.
• Figure 2: Comparative Demonstration of the Stream Aligner Module. Typically, the Stream Aligner has two working patterns depending on the original suffix sentence: (i) If the suffix is correct, copy it; (ii) If the suffix is incorrect or unsatisfying, rewrite it to make it correct and better. In this way, Stream Aligner can eliminate minor mistakes and toxic output made by the upstream model, thereby eliciting more correct latent knowledge and reducing the reliance on Stream Aligner capability.
• Figure 3: Performance of Stream Aligner models. (a) The win rate of Llama2-70B-chat + Stream Aligner-2B on helpfulness and harmlessness, compared to the baseline generated by Llama2-70B-chat. (b) The win rate(correct rate) of Llama3-8B-Instruct + Stream Aligner-2B and Llama3-70B-Instruct + Stream Aligner-8B on math, compared to the zero-shot baseline of Llama3-8B-Instruct and Llama3-70B-Instruct. It is demonstrated that Stream Aligner achieves excellent performances across all evaluation metrics for every task. Furthermore, the overall performance of Stream Aligner tends to increase with the number of correction rounds, eventually converging to a stable value. It is also observable that, compared to helpful & harmless QA tasks, the performance of Stream Aligner on math tasks almost monotonically increases with each round.
• Figure 4: Distribution of helpful and harmless scores across different rounds.(a-c) show the distribution of helpful and harmless scores during the evaluation generation process. From the development of the distribution, we can find: (1) In the first few rounds, since the output length is small, the helpfulness score is relatively low. In the middle and last rounds, the answers can be positively corrected to be helpful (2) The harmlessness score of each round continues to increase during the correction generation process. This might because the corrected suffix of Stream Aligner is generally safe, and the more corrected output is generated, the higher the safety score would be. (3) Overall, we can see the Stream Aligner tends to correct the original answer into another distribution, which is more helpful and harmless than the original distribution.
• Figure 5: Ablation of generation pipeline. We performed an ablation study about different generation pipelines on math tasks and different sizes of upstream models. We observe that under different upstream and Stream Aligner models, the new continue generation pipeline exceeds the classical sentence-by-sentence correction pipeline in terms of performance on math tasks. When correction rounds increase, the win rate of both pipelines will eventually converge to a constant value.