PaperlayerView exact revision on arXiv

Source-linked AI summary

From Drafts to Answers: Unlocking LLM Potential via Aggregation Fine-Tuning

Yafu Li, Zhilin Wang, Tingchen Fu, Ganqu Cui, Sen Yang, Yu Cheng

arXiv:2501.11877v1cs.CLcs.AI

TL;DR

Inference-time computation can improve LLM performance, but intermediate processes for general tasks remain less discussed. AFT trains models to aggregate multiple drafts into refined answers, and propose-and-aggregate further improves performance; the top reported AFT model reached a 41.3% LC win rate on AlpacaEval 2, surpassing much larger models.

  • Problem

    Although inference-time computation can boost performance, intermediate processes for general tasks such as writing and open-ended dialogue remain less discussed than those for coding and math.

    Source: §1 · p. 1

  • Method

    AFT trains a model to aggregate multiple proposals into a refined answer, and propose-and-aggregate iteratively generates and aggregates proposals at inference time.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    AFT substantially outperforms standard SFT across benchmarks; the top AFT-on-policy Llama3.1-8B-Base model achieved a 41.3% LC win rate on AlpacaEval 2, surpassing GPT-4 and Llama3.1-405B-Instruct.

    Source: §1 · p. 2, §5.1 · p. 6

  • Takeaways & Limitations

    Propose-and-aggregate combines sequential refinement and parallel sampling for inference-time scaling, with consistent multi-task improvements and manageable computational overhead.

    Source: §1 · p. 2, §7 · p. 9

Abstract

from arXiv · show

Scaling data and model size has been proven effective for boosting the performance of large language models. In addition to training-time scaling, recent studies have revealed that increasing test-time computational resources can further improve performance. In this work, we introduce Aggregation Fine-Tuning (AFT), a supervised finetuning paradigm where the model learns to synthesize multiple draft responses, referred to as proposals, into a single, refined answer, termed aggregation. At inference time, a propose-and-aggregate strategy further boosts performance by iteratively generating proposals and aggregating them. Empirical evaluations on benchmark datasets show that AFT-trained models substantially outperform standard SFT. Notably, an AFT model, fine-tuned from Llama3.1-8B-Base with only 64k data, achieves a 41.3% LC win rate on AlpacaEval 2, surpassing significantly larger LLMs such as Llama3.1-405B-Instruct and GPT4. By combining sequential refinement and parallel sampling, the propose-and-aggregate framework scales inference-time computation in a flexible manner. Overall, These findings position AFT as a promising approach to unlocking additional capabilities of LLMs without resorting to increasing data volume or model size.

1. Introduction

The introduction argues that inference-time computation can improve model performance, while intermediate processes for general tasks remain less explored. It presents AFT, which teaches models to aggregate drafts, and reports gains over standard SFT and further improvements from propose-and-aggregate inference.

1. Introduction

  • Although inference-time computation can boost performance, intermediate processes for general tasks such as writing and open-ended dialogue remain less discussed than those for coding and math.Typical SFT still maps a user query directly to a reference response.Source: §1 · p. 1
  • AFT trains a model to synthesize multiple draft responses into a refined answer, rather than mapping a query directly to a response.Its training examples contain a query, proposals, and an aggregation; proposals may be off-policy or on-policy.Source: Figure 1 · p. 1, §1 · p. 2
  • 41.3% LC win rate on AlpacaEval 2 was achieved by the top AFT model, which surpassed GPT4 and Llama3.1-405B-Instruct after fine-tuning on 64k examples.The introduction also reports that AFT improves LC win rates by up to 13.9% on average and that propose-and-aggregate raises Mistral-based models above 30% and Llama-based models above 40%.Source: §1 · p. 2
  • Analysis finds that conditioning on drafts shapes a low-perplexity region associated with faster, more stable convergence, while proposal diversity and quality both contribute to aggregation quality.Source: §1 · p. 2
  • At inference, propose-and-aggregate iteratively generates proposals and aggregates them, combining parallel sampling with sequential refinement to improve performance across tasks.The introduction describes this scaling approach as having manageable computational overhead.Source: §1 · p. 2

2. Related Work

Related work covers data and intermediate-process approaches to SFT, alongside methods for scaling computation at inference. This paper adapts aggregation from mixture-of-agents to a single-model setting for training and inference.

Supervised Fine-tuning of LLMs

  • Prior SFT work studies instruction-data quality and curation, including constrained, multifaceted, self-refined, and sequential data.The passage also contrasts aggregation fine-tuning as using aggregation for general-purpose instruction following.Source: §2 · p. 2

Inference Scaling for LLMs

  • Inference-time scaling includes sequential revision, parallel sampling, and step-level graph search; prior methods include both single-model algorithms and mixture-of-agents.The passage notes that smaller models can surpass larger ones under a fixed test-time budget.Source: §2 · pp. 2–3
  • This work brings mixture-of-agents aggregation into a single-model setting, using aggregation for general instruction following and combining sequential revision with parallel sampling at test time.Source: §2 · p. 3

3. Method

Aggregation fine-tuning trains a model to refine multiple proposals into a reference answer, rather than map a query directly to one response. Its data construction, training objective, and iterative inference extend this aggregation process from supervision to test-time computation.

3. Method

  • AFT trains an aggregator to turn a query and K proposals into a refined response approximating a high-quality reference, unlike direct query-to-response fine-tuning.The model is trained to evaluate, critique, and improve initial solutions by conditioning on the drafts.Source: §3 · p. 3

3.1. Data Construction

  • Off-policy proposals come from other models' responses in preference datasets, while on-policy proposals are sampled from the model being fine-tuned; a stronger model creates the reference aggregation.Each training instance contains a query, proposals, and a refined reference response.Source: §3.1 · p. 3

3.2. Training

  • Training maximizes the aggregator's probability of producing the reference aggregation from the query and proposals, whose semantic cues and partial solutions reduce generation perplexity.The proposal-conditioned learning signal is described as producing a more stable and smooth training curve.Source: §3.2 · p. 3

3.3. Inference

  • At inference, the model samples initial proposals, aggregates them, and can iteratively use each layer's outputs as proposals for the next, combining sequential refinement with parallel sampling.Iteration continues to a predefined layer limit; Figure 3 describes scaling along depth and breadth.Source: §3.3 · p. 4

4. Experimental Setup

Experiments use UltraFeedback-derived data to train Llama and Mistral base models with LoRA, then evaluate dialogue and downstream-task performance. Inference settings vary by benchmark, and the setup reports AlpacaEval 2 results in Table 1.

Data

  • UltraFeedback's 64k queries and model responses provide off-policy proposals; on-policy data uses ten sampled responses per query, with a reward model selecting five proposals before Qwen2.5-72B-Instruct generates aggregations.The response sampling temperature is 1.0, with top k 50 and top p 0.95.Source: §4 · pp. 4–5

Training

  • Llama-3.1-8B-Base and Mistral-7B-v0.1-Base models are fine-tuned with LoRA for one epoch, using rank 8, batch size 128, and a 1 × 10^-4 learning rate.Training runs across four A100 GPUs with a cosine schedule and 10% warmup.Source: §4 · p. 5

Evaluation

  • Evaluation covers AlpacaEval 2 and MT-Bench dialogue benchmarks plus MMLU, ARC-c, StrategyQA, GSM8K, and IFEval downstream tasks; Table 1 reports AlpacaEval 2 single-turn results.AlpacaEval 2 reports raw and length-controlled win rates; MT-Bench reports an average score judged by GPT-4.Source: §4 · p. 5, Table 1 · p. 5

Inference

  • Inference defaults to five proposals and two aggregation layers for AlpacaEval 2 and MT-Bench, but uses one layer for GSM8K and IFEval.Proposal generation and intermediate aggregation use temperature 0.7 and top p 0.95.Source: §4 · p. 5

5. Experimental Results

The experiments compare aggregation fine-tuning (AFT) with standard supervised fine-tuning across dialogue and downstream-task benchmarks. They also assess whether propose-and-aggregate inference further improves performance.

5.1. Main Results

  • Without inference scaling, AFT improves average AlpacaEval 2 LC win rate by 14.9% for Mistral-based models and 12.9% for Llama-based models over SFT.AFT-on-policy outperforms AFT-off-policy in both model families.Source: §5.1 · p. 5
  • 8.1 average MT-Bench score for AFT-on-policy Llama3.1-8B-Base with propose-and-aggregate exceeds 7.4 without it and 6.8 for SFT.Table 2 reports multi-turn dialogue evaluation; Figure 4 presents a radar chart for Llama3.1-8B-Base models.Source: Table 2 · p. 6, §5.1 · p. 6, Figure 4 · p. 6
  • 41.3% LC win rate on AlpacaEval 2 for the top AFT-on-policy Llama3.1-8B-Base model, fine-tuned on 64k examples, surpasses GPT-4 and Llama3.1-405B-Instruct.The result uses no further preference alignment.Source: §5.1 · p. 6

5.2. Downstream Task Performance

  • AFT matches SFT on MMLU while improving ARC-c and StrategyQA; on GSM8K and IFEval, AFT outperforms SFT, with further gains from propose-and-aggregate.For these tasks, SFT models cannot use propose-and-aggregate to boost performance. Figure 5 summarizes GSM8K and IFEval results.Source: §5.2 · p. 6, Table 3 · p. 6, Figure 5 · p. 7

6. Analysis

Section 6 analyzes aggregation learning, proposal quality and diversity, test-time scaling across width and depth, and computational overhead.

6.1. Understanding Aggregation Learning

  • AFT trains more efficiently and stably than SFT, while removing proposals or substituting unaggregated pseudo-targets raises perplexity and degrades performance.AFT data has lower perplexity than SFT data, and on-policy drafts further reduce it; the authors interpret this as mode-seeking.Source: §6.1 · p. 7, Table 4 · p. 7

6.2. Effects of Proposal Diversity and Quality

  • Higher proposal quality and diversity generally improve aggregation quality, with proposal quality having the stronger effect.Figure 7 encodes aggregation quality in grid cells, with lighter colors indicating higher quality.Source: Figure 7 · p. 8, §6.2 · p. 8

6.3. Test-time Scaling along Width and Depth

  • Increasing proposal count and aggregation depth improves performance, while one proposal or one aggregation layer reduces the framework to sequential revision or parallel sampling, respectively.With one aggregation layer, the model itself effectively acts as a verifier through aggregation.Source: Figure 8 · p. 8, §6.3 · p. 8

6.4. Computational Overhead

  • Even with BoN, SFT cannot surpass AFT, and propose-and-aggregate improves AFT more than BoN without an external reward model.Its additional FLOPs arise from aggregation layers processing proposals from the previous layer; the approximation is ˆF ≈ L · (2 · N · F) + F.Source: §6.4 · pp. 8–9, Table 5 · p. 8

7. Conclusion

AFT trains language models to refine multiple proposals into a final answer and shows consistent gains over standard SFT. Iterative propose-and-aggregate improves performance without additional training, while analysis highlights effects on perplexity, convergence, and proposal balance.

  • AFT demonstrates consistent gains over standard supervised fine-tuning across instruction-following benchmarks and downstream tasks.Source: §7 · p. 9
  • Iterative propose-and-aggregate further improves performance without additional training by combining sequential refinement and parallel sampling.Source: §7 · p. 9
  • AFT teaches models to combine multiple draft responses, or proposals, into a refined final answer.Source: §7 · p. 9
  • Analysis finds reduced perplexity, stabilized convergence, and benefits from balancing proposal quality with diversity.Source: §7 · p. 9

Impact Statement

The statement presents Aggregation Fine-Tuning as a training and inference paradigm for refining candidate answers. It notes possible misuse risks and potential benefits from more accurate, reliable, and resource-efficient language models.

  • Aggregation Fine-Tuning is presented as a paradigm that refines multiple candidate answers into a stronger overall solution.Source: Impact Statement · p. 9
  • The statement identifies possible misuse of more powerful language generation for deceptive or harmful purposes, such as misleading content.Source: Impact Statement · p. 9
  • It encourages robust filtering and responsible-use guidelines, and notes potential benefits in applications including education and healthcare.Source: Impact Statement · p. 9

A. Prompt

Aggregation uses the same prompt as Mixture-of-Agents, with proposals embedded in a template that is input to the model as a system prompt.

  • The aggregation prompt is the same one used by Mixture-of-Agents.Source: A. Prompt · p. 15
  • Proposals are embedded within the aggregation template.Source: A. Prompt · p. 15
  • The template containing the proposals is input to the LLM as a system prompt.Source: A. Prompt · p. 15

B. Case Study

The case studies illustrate propose-and-aggregate generation across mathematics, reasoning, knowledge, writing, and role-play tasks.

C. Computational Overhead

The section estimates inference FLOPs for propose-and-aggregate and parallel sampling, alongside per-token and total-inference estimates.

C. Computational Overhead

  • FLOPs per token ≈ 2 × (num parameters + 2 × num layers × token dim × context length), and total inference FLOPs ≈ (num prompt tokens + num decoded tokens) × FLOPs per token.The text introduces these as approximations for per-token and total inference FLOPs.Source: C. Computational Overhead · p. 15
  • For simplicity, the analysis omits query-prompt FLOPs and focuses on overhead from prompt tokens accommodating previous-layer proposals.This is the assumption used for the propose-and-aggregate overhead estimate.Source: C. Computational Overhead · p. 15
  • For parallel sampling with a reward model, ¯F ≈ 2 · num proposals · num decoded tokens · FLOPs per token, with the factor 2 accounting for RM evaluation.The RM is typically comparable in size to the policy model.Source: C. Computational Overhead · p. 15
  • For propose-and-aggregate, ˆF ≈ (num aggregation layers) · (2 · num proposals · num decoded tokens · FLOPs per token) + (num decoded tokens · FLOPs per token).The first term covers proposal generation and processing across aggregation layers; the second covers final generation.Source: C. Computational Overhead · pp. 15–16
Loading 2501.11877v1…