Logits are All We Need to Adapt Closed Models

2025-02-03

AI Summary

Closed-source large language models (LLMs) cannot expose internal parameters or gradients, limiting adaptation to prompt engineering alone. Method: We propose a lightweight black-box task alignment framework that reformulates closed-source LLM adaptation as a label-noise correction problem. Theoretically, we prove that reliable task transfer can be achieved via probabilistic reweighting of output token logits—requiring no parameter access, gradient computation, or architectural modification. Our plug-and-play logits reweighting module integrates supervised classification–guided reweighting, label-noise-robust learning, and autoregressive probability calibration, enabling efficient fine-tuning with minimal task data. Contribution/Results: Experiments across multiple closed-source LLMs (e.g., GPT-4, Claude) and diverse benchmark datasets demonstrate significant improvements over prompt tuning baselines. Our approach validates that controllable, logits-level intervention substantially enhances the downstream utility and adaptability of closed-source models.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Many commercial Large Language Models (LLMs) are often closed-source, limiting developers to prompt tuning for aligning content generation with specific applications. While these models currently do not provide access to token logits, we argue that if such access were available, it would enable more powerful adaptation techniques beyond prompt engineering. In this paper, we propose a token-level probability reweighting framework that, given access to logits and a small amount of task-specific data, can effectively steer black-box LLMs toward application-specific content generation. Our approach views next-token prediction through the lens of supervised classification. We show that aligning black-box LLMs with task-specific data can be formulated as a label noise correction problem, leading to emph{Plugin} model -- an autoregressive probability reweighting model that operates solely on logits. We provide theoretical justification for why reweighting logits alone is sufficient for task adaptation. Extensive experiments with multiple datasets, LLMs, and reweighting models demonstrate the effectiveness of our method, advocating for broader access to token logits in closed-source models.

Problem

• Adapt closed-source Large Language Models
• Token-level probability reweighting framework
• Align LLMs with task-specific data

Innovation

• Token-level probability reweighting framework
• Supervised classification for next-token prediction
• Autoregressive probability reweighting on logits