Add Sequential Temperature Scaling for confidence-head calibration

deepspec/eval/dspark/confidence_head.py

@@ -322,6 +322,7 @@ def __init__(
         tensorboard_dir: str | None,
         step: int | None,
         artifact_root: Path | None,
+        records_dir: Path | None = None,
     ):
         self.device = device
         self.max_proposal_tokens = int(max_proposal_tokens)
@@ -331,7 +332,9 @@ def __init__(
         self.tensorboard_dir = tensorboard_dir
         self.step = step
         self.artifact_root = artifact_root
+        self.records_dir = records_dir
         self.dataset_metrics: PerPositionConfidenceMetrics | None = None
+        self.records: list[dict] = []
         self.rows: list[dict] = []
 
     def start(self) -> None:
@@ -341,6 +344,7 @@ def start(self) -> None:
             num_fine_bins=self.num_fine_bins,
             device=self.device,
         )
+        self.records = []
 
     def observe(
         self,
@@ -373,6 +377,39 @@ def observe(
             probs=cumprod_pred,
             targets=prefix_label,
         )
+        if self.records_dir is not None:
+            self.records.append(
+                {
+                    "logits": confidence_logits[0, :effective_length]
+                    .detach()
+                    .to(torch.float32)
+                    .cpu()
+                    .tolist(),
+                    "labels": prefix_label.long().cpu().tolist(),
+                }
+            )
+
+    def _gather_records(self) -> list[dict]:
+        records = self.records
+        self.records = []
+        if self.records_dir is None:
+            return []
+        gathered: list[list[dict] | None] = [None] * dist.get_world_size()
+        dist.all_gather_object(gathered, records)
+        merged: list[dict] = []
+        for rank_records in gathered:
+            merged.extend(rank_records or [])
+        return merged
+
+    def _write_records(self, *, dataset_name: str, records: list[dict]) -> Path:
+        assert self.records_dir is not None
+        dataset_dir = self.records_dir / dataset_name
+        dataset_dir.mkdir(parents=True, exist_ok=True)
+        records_path = dataset_dir / "confidence_records.jsonl"
+        with records_path.open("w", encoding="utf-8") as handle:
+            for record in records:
+                handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+        return records_path
 
     def finish(
         self,
@@ -384,10 +421,18 @@ def finish(
         dataset_metrics = self.dataset_metrics
         dataset_metrics.all_reduce()
         self.dataset_metrics = None
+        records = self._gather_records()
 
         if dist.get_rank() != 0 or int(metric_summary["sample_count"]) == 0:
             return None
 
+        if self.records_dir is not None and records:
+            records_path = self._write_records(
+                dataset_name=dataset_name,
+                records=records,
+            )
+            print(f"Wrote confidence records to {records_path}", flush=True)
+
         row = self.build_dataset_row(
             dataset_name=dataset_name,
             metric_summary=metric_summary,

deepspec/eval/dspark/draft_ops.py

@@ -101,6 +101,7 @@ def build_dspark_proposal(
     block_size: int,
     temperature: float,
     confidence_threshold: float,
+    confidence_temperatures: torch.Tensor | None = None,
 ) -> DSparkDraftProposal:
     assert draft_input_ids.size(0) == 1, "build_dspark_proposal requires batch_size=1"
     proposal_hidden_states = block_hidden[:, :block_size, :]
@@ -124,6 +125,12 @@ def build_dspark_proposal(
         )
         if confidence_logits is None:
             return _empty_dspark_proposal(draft_input_ids)
+        if confidence_temperatures is not None:
+            # Sequential Temperature Scaling: divide each position's logit by
+            # its fitted temperature so downstream consumers (the threshold
+            # early-stop and the calibration metrics) see calibrated
+            # probabilities.
+            confidence_logits = confidence_logits / confidence_temperatures
         proposal_draft_tokens = _confident_prefix_length(
             confidence_logits,
             block_size=block_size,

deepspec/eval/dspark/evaluator.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 from pathlib import Path
 from types import SimpleNamespace
 
@@ -35,12 +36,37 @@ class Qwen3DSparkEvaluator(BaseEvaluator):
 
     def __init__(self, local_rank: int, args):
         super().__init__(local_rank, args)
+        self.confidence_temperatures = self._load_confidence_temperatures()
         self.confidence_head_recorder = self._build_confidence_head_recorder()
 
     @property
     def max_proposal_tokens(self) -> int:
         return int(self.draft_model.block_size)
 
+    def _load_confidence_temperatures(self) -> torch.Tensor | None:
+        path = getattr(self.args, "confidence_calibration_path", None)
+        if path is None:
+            return None
+        assert self.draft_model.confidence_head is not None, (
+            "--confidence-calibration-path requires a draft model with a "
+            "confidence head."
+        )
+        with Path(path).open("r", encoding="utf-8") as handle:
+            payload = json.load(handle)
+        temperatures = [float(value) for value in payload["temperatures"]]
+        assert len(temperatures) == self.max_proposal_tokens, (
+            f"Calibration file {path} has {len(temperatures)} temperatures, "
+            f"expected block_size={self.max_proposal_tokens}."
+        )
+        assert all(value > 0.0 for value in temperatures), (
+            f"Calibration temperatures must be positive, got {temperatures}."
+        )
+        return torch.tensor(
+            temperatures,
+            dtype=torch.float32,
+            device=self.device,
+        )
+
     def _build_confidence_head_recorder(self) -> ConfidenceHeadRecorder | None:
         if self.draft_model.confidence_head is None:
             return None
@@ -54,6 +80,9 @@ def _build_confidence_head_recorder(self) -> ConfidenceHeadRecorder | None:
                 / "artifacts"
                 / f"step_{self.args.step}"
             )
+        records_dir = None
+        if getattr(self.args, "confidence_dump_dir", None) is not None:
+            records_dir = Path(self.args.confidence_dump_dir)
         return ConfidenceHeadRecorder(
             device=self.device,
             max_proposal_tokens=self.max_proposal_tokens,
@@ -63,6 +92,7 @@ def _build_confidence_head_recorder(self) -> ConfidenceHeadRecorder | None:
             tensorboard_dir=self.args.tensorboard_dir,
             step=self.args.step,
             artifact_root=artifact_root,
+            records_dir=records_dir,
         )
 
     def build_models(self) -> tuple[object, Qwen3DSparkModel, AutoTokenizer]:
@@ -129,6 +159,7 @@ def _propose(
             block_size=self.max_proposal_tokens,
             temperature=float(self.args.temperature),
             confidence_threshold=float(self.args.confidence_threshold),
+            confidence_temperatures=self.confidence_temperatures,
         )
 
     def _update(

eval.py

@@ -39,6 +39,27 @@ def parse_args():
         default=0.0,
         help=("Confidence-head early-stop threshold. Confidence calibration metrics are collected only when this is 0.0."),
     )
+    parser.add_argument(
+        "--confidence-dump-dir",
+        type=str,
+        default=None,
+        help=(
+            "Dump per-proposal confidence records to "
+            "<dir>/<dataset>/confidence_records.jsonl for offline calibration "
+            "fitting (see scripts/fit_confidence_calibration.py). Requires "
+            "--confidence-threshold 0."
+        ),
+    )
+    parser.add_argument(
+        "--confidence-calibration-path",
+        type=str,
+        default=None,
+        help=(
+            "JSON file with per-position confidence temperatures fitted by "
+            "scripts/fit_confidence_calibration.py; the temperatures are "
+            "applied to the confidence logits during evaluation."
+        ),
+    )
     parser.add_argument("--tensorboard-dir", type=str, default=None)
     parser.add_argument("--step", type=int, default=None,help=("step for tensorboard logging"),)
     parser.add_argument("--seed", type=int, default=980406)