Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unlocking the Pre-Trained Model as a Dual-Alignment Calibrator for Post-Trained LLMs

Beier Luo, Cheng Wang, Hongxin Wei, Sharon Li, Xuefeng Du

TL;DR

This work addresses miscalibration in post-trained LLMs by revealing two drift mechanisms—confidence drift and process drift—that underlie overconfident predictions. It introduces Dual-Align, a post-hoc, unsupervised calibration framework that combines confidence alignment with process alignment, using a Peak Divergence Layer to trigger ISE-based stabilization and a learned temperature $\tau$ to balance objectives. The method unifies final-distribution matching with trajectory-level stabilization, guided by a per-sample divergence weight, and demonstrates substantial reductions in calibration errors across multiple LLMs and benchmarks, approaching supervised oracle performance. The approach is robust across open-ended tasks and various post-training methods, suggesting broad applicability for reliable, calibrated inference in PoLMs without additional labeled data.

Abstract

Post-training improves large language models (LLMs) but often worsens confidence calibration, leading to systematic overconfidence. Recent unsupervised post-hoc methods for post-trained LMs (PoLMs) mitigate this by aligning PoLM confidence to that of well-calibrated pre-trained counterparts. However, framing calibration as static output-distribution matching overlooks the inference-time dynamics introduced by post-training. In particular, we show that calibration errors arise from two regimes: (i) confidence drift, where final confidence inflates despite largely consistent intermediate decision processes, and (ii) process drift, where intermediate inference pathways diverge. Guided by this diagnosis, we propose Dual-Align, an unsupervised post-hoc framework for dual alignment in confidence calibration. Dual-Align performs confidence alignment to correct confidence drift via final-distribution matching, and introduces process alignment to address process drift by locating the layer where trajectories diverge and realigning the stability of subsequent inference. This dual strategy learns a single temperature parameter that corrects both drift types without sacrificing post-training performance gains. Experiments show consistent improvements over baselines, reducing calibration errors and approaching a supervised oracle.

Unlocking the Pre-Trained Model as a Dual-Alignment Calibrator for Post-Trained LLMs

TL;DR

This work addresses miscalibration in post-trained LLMs by revealing two drift mechanisms—confidence drift and process drift—that underlie overconfident predictions. It introduces Dual-Align, a post-hoc, unsupervised calibration framework that combines confidence alignment with process alignment, using a Peak Divergence Layer to trigger ISE-based stabilization and a learned temperature to balance objectives. The method unifies final-distribution matching with trajectory-level stabilization, guided by a per-sample divergence weight, and demonstrates substantial reductions in calibration errors across multiple LLMs and benchmarks, approaching supervised oracle performance. The approach is robust across open-ended tasks and various post-training methods, suggesting broad applicability for reliable, calibrated inference in PoLMs without additional labeled data.

Abstract

Post-training improves large language models (LLMs) but often worsens confidence calibration, leading to systematic overconfidence. Recent unsupervised post-hoc methods for post-trained LMs (PoLMs) mitigate this by aligning PoLM confidence to that of well-calibrated pre-trained counterparts. However, framing calibration as static output-distribution matching overlooks the inference-time dynamics introduced by post-training. In particular, we show that calibration errors arise from two regimes: (i) confidence drift, where final confidence inflates despite largely consistent intermediate decision processes, and (ii) process drift, where intermediate inference pathways diverge. Guided by this diagnosis, we propose Dual-Align, an unsupervised post-hoc framework for dual alignment in confidence calibration. Dual-Align performs confidence alignment to correct confidence drift via final-distribution matching, and introduces process alignment to address process drift by locating the layer where trajectories diverge and realigning the stability of subsequent inference. This dual strategy learns a single temperature parameter that corrects both drift types without sacrificing post-training performance gains. Experiments show consistent improvements over baselines, reducing calibration errors and approaching a supervised oracle.
Paper Structure (47 sections, 2 theorems, 28 equations, 12 figures, 7 tables)

This paper contains 47 sections, 2 theorems, 28 equations, 12 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Our contributions.
  3. Preliminaries
  4. Probability distribution across transformer layers.
  5. Confidence calibration for PoLMs.
  6. Post-hoc calibration methods.
  7. Understanding the Effects of Post-training
  8. Confidence drift.
  9. Process drift.
  10. Methodology
  11. Stability of Inference Across Transformer Layers
  12. Process Alignment for Process Drift
  13. Dual-Align: A Unified Calibration Framework
  14. Experiments
  15. Experimental Setup
  16. ...and 32 more sections

Key Result

Proposition 1

(Informal). Under mild regularity conditions (Appendix app:theory), there exist bounded weights $w(\bm{x})\ge 0$ such that where $\mathrm{ISE}_f(\bm{x},\tau)$ is the PoLM inferential stability entropy under temperature $\tau$, $\mathrm{ISE}_g(\bm{x})$ is the PLM stability reference, and $C_g$ is a positive constant relevant to the PLM.

Figures (12)

  • Figure 1: Illustration of our method: DUAL-ALIGN. Our approach addresses both confidence drift and process drift. For confidence drift, we align the LLMs’ confidence using the objective $\mathcal{L}_{\rm Conf}$ (Left). For process drift, we first identify the Peak Divergence Layer (PDL), then calculate the Inferential Stability Entropy (ISE) with respect to the process drift between the PLM and PoLM, and align it using the objective $\mathcal{L}_{\rm Process}$ (Right).
  • Figure 2: The layer-wise Jensen-Shannon Divergence between a post-trained mode Llama-3.1-8B-Instruct and a pre-trained model Llama-3.1-8B on MMLU. Agreed samples show minimal differences, suggesting confidence drift, while disagreed samples display a sharp spike at an intermediate layer, indicating process drift.
  • Figure 3: Relationship between output confidence and Inferential Stability Entropy (ISE) of Qwen2.5-14B nad Qwen2.5-14B-Instruct on MMLU. The well-calibrated pre-trained model (left) displays an ISE distribution similar to a normal distribution, whereas the post-trained model (right) shows extreme overconfidence and abnormally low ISE values, indicating overly rigid decision-making processes.
  • Figure 4: (a) Applicability to open-ended question answering. We evaluate LLama3.1 and Qwen2.5-14B on TruthfulQA dataset. (b) Applicability to different post-training methods. Apart from instruction-tuning, we consider PPO, DPO and GRPO training on Qwen2.5-7B.
  • Figure 5: Reliability diagrams on MMLU comparing a PLM with PoLMs obtained through various post-training methods. The pre-trained model is Llama-3.1-8B-Base and the post-training techniques include Supervised Fine-tuning (SFT), Direct Preference Optimization (DPO) and Group Relative Policy Optimization (GRPO).
  • ...and 7 more figures

Theorems & Definitions (3)

  • Proposition 1
  • Theorem 1: Formal version of Proposition \ref{['prop:process_align_calibration']}
  • Proof 1