Research map / Large language models / Language, safety and interpretability
Logits are All We Need to Adapt Closed Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Plugin, an autoregressive token-level probability reweighting model that adapts black-box LLMs using only output logits and a small target dataset, framed as label-noise correction with a diagonal transition matrix.
Problem
Adapting closed-source LLMs to application-specific content generation when only token logits (not weights or training data) are accessible, where prompt tuning is insufficient.
Approach
Views next-token prediction as supervised classification with label noise; estimates a diagonal autoregressive transition vector via a small reweighting model trained on target data, and multiplies its probabilities with the black-box model's softmax output (normalized elementwise product) at inference.
Key results
- On E2E NLG with GPT2-XL, Plugin reaches BLEU 0.2470 and ROUGE-L 0.4213 vs 0.1325 and 0.3021 for TempNet (best baseline shown in table for that model).
- On Web NLG with LLaMA-3.1-8B, Plugin improves person-related generation from 17.99% to 71.34% (GPT-4o judged) over the fine-tuned base model.
- Human evaluation on 100 Adidas samples prefers Plugin over ICL-3 in 81% of cases.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Plugin instead of the strongest reported baseline for BLEU | GPT2-M, GPT2-XL, or LLaMA-3.1-8B · E2E NLG test | BLEU: 0.2377$_{\pm0.011}$ -> 0.1863$_{\pm0.010}$; 0.0980$_{\pm0.035}$ -> 0.2470$_{\pm0.009}$; 0.3527$_{\pm0.033}$ -> 0.3691$_{ | Mixed |
| Use Plugin instead of the strongest reported baseline for BLEU | GPT2-M, GPT2-XL, or LLaMA-3.1-8B · Web NLG test | BLEU: 0.1071$_{\pm0.005}$ -> 0.1280$_{\pm0.007}$; 0.1071$_{\pm0.005}$ -> 0.1673$_{\pm0.004}$; 0.2315$_{\pm0.010}$ -> 0.2542$_{ | Helped |
| Use Plugin instead of the strongest reported baseline for BLEU | GPT2-M, GPT2-XL, or LLaMA-3.1-8B · CommonGen test | BLEU: 0.1260$_{\pm0.007}$ -> 0.1366$_{\pm0.003}$; 0.1260$_{\pm0.007}$ -> 0.1791$_{\pm0.014}$; 0.1789$_{\pm0.005}$ -> 0.2665$_{ | Helped |
| Use Plugin instead of the strongest reported baseline for BLEU | GPT2-M, GPT2-XL, or LLaMA-3.1-8B · Adidas test | BLEU: 0.0565$_{\pm0.014}$ -> 0.0486$_{\pm0.006}$; 0.0567$_{\pm0.016}$ -> 0.0600$_{\pm0.017}$; 0.0506$_{\pm0.011}$ -> 0.0611$_{ | Mixed |
| Add Plugin on top of BDPL | BDPL with GPT2-XL · E2E NLG | BLEU: 0.2287 -> 0.4527 | Helped |
| Add Plugin on top of BDPL | BDPL with GPT2-XL · Web NLG | BLEU: 0.1024 -> 0.2137 | Helped |
| Add Plugin on top of BDPL | BDPL with GPT2-XL · CommonGen | BLEU: 0.1023 -> 0.2614 | Helped |
| Add Plugin on top of BDPL | BDPL with GPT2-XL · Adidas | BLEU: 0.0417 -> 0.0623 | Helped |
| Increase the base model's task-specific fine-tuning (ablation) | Plugin with GPT2-M · E2E NLG | BLEU and ROUGE-L | Helped |
| Increase reweighting model depth beyond one transformer layer (ablation) | Plugin with GPT2-M · E2E NLG | BLEU and ROUGE-L | Hurt |
| Initialize the reweighting model with GPT2-Small (ablation) | Plugin with GPT2-M · E2E NLG | BLEU and ROUGE-L | Helped |
| Use Plugin rather than the fine-tuned base model | LLaMA-3.1-8B · Web NLG distribution shift | Person-related generated sentences: 17.99% -> 71.34% | Helped |
Surprising findings
- NewModel achieves higher NIST scores on some datasets, attributed to repetition of less-frequent input tokens, despite lower coherence on other metrics.
- Adding more layers to the reweighting model degrades performance on E2E NLG due to overfitting on the small validation set.
Limitations
- Only reweights token probabilities without modifying internal representations, so may struggle with deep structural adaptation such as complex reasoning.
- Training a separate reweighting model adds computational overhead relative to prompt tuning or in-context learning.
- Experiments use greedy decoding, so absolute numbers are not competitive with state-of-the-art; diagonal transition assumption is a restrictive simplification.
Experimental setup
Datasets: E2E NLG Web NLG CommonGen Adidas product descriptions
Benchmarks: E2E NLG Web NLG CommonGen Adidas
Metrics: BLEU ROUGE-1 ROUGE-2 ROUGE-L METEOR CIDEr NIST
Baselines: Zeroshot ICL-1 ICL-3 NewModel WeightedComb TempNet BDPL
Builds on: forward loss correction (label noise) plugin classifiers under distribution shift DExperts BDPL
Closest papers by meaning
- Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining (ICLR 2025 · similarity 0.88)
- Rethinking Data Curation in LLM Training: Online Reweighting Offers Better Generalization than Offline Methods (ICLR 2026 · similarity 0.88)
- Reinforced Fast Weights via Next-Sequence Prediction (NeurIPS 2026 · similarity 0.88)
- Rethinking Personalization in Large Language Models at the Token Level (ICML 2026 · similarity 0.88)
- Better & Faster Large Language Models via Multi-token Prediction (ICML 2024 · similarity 0.87)
- Exploring Criteria of Loss Reweighting to Enhance LLM Unlearning (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
