蒸馏把温度升高，为什么还要把损失乘回 T²？

"""Exact arithmetic on fixed teacher/student logits; no neural network training."""
import json
import math
from pathlib import Path


def log_softmax(logits, temperature):
    if not logits or not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("nonempty logits and positive finite temperature required")
    if any(not math.isfinite(x) for x in logits):
        raise ValueError("finite logits required")
    scaled = [x / temperature for x in logits]
    high = max(scaled)
    log_sum = math.log(sum(math.exp(x-high) for x in scaled))
    return [x-high-log_sum for x in scaled]


def kl_and_gradient(teacher, student, temperature, rescale):
    if len(teacher) != len(student):
        raise ValueError("same ordered classes required")
    lp = log_softmax(teacher, temperature)
    lq = log_softmax(student, temperature)
    p, q = [math.exp(v) for v in lp], [math.exp(v) for v in lq]
    multiplier = temperature ** 2 if rescale else 1.0
    loss = multiplier * sum(pi*(a-b) for pi, a, b in zip(p, lp, lq))
    gradient = [multiplier*(qi-pi)/temperature for pi, qi in zip(p, q)]
    return loss, gradient, p, q


def norm(values):
    return math.sqrt(sum(v*v for v in values))


def main():
    teacher, student = [4.0, 1.0, -2.0], [1.0, 0.0, -1.0]
    results, errors = [], []
    for temperature in [1.0, 2.0, 4.0, 16.0]:
        raw, raw_grad, p, q = kl_and_gradient(teacher, student, temperature, False)
        scaled, grad, _, _ = kl_and_gradient(teacher, student, temperature, True)
        eps = 1e-4
        for j in range(len(student)):
            plus, minus = student.copy(), student.copy()
            plus[j] += eps
            minus[j] -= eps
            numeric = (kl_and_gradient(teacher, plus, temperature, True)[0]
                       - kl_and_gradient(teacher, minus, temperature, True)[0]) / (2*eps)
            errors.append(abs(numeric-grad[j]))
        assert math.isclose(sum(grad), 0.0, abs_tol=1e-12)
        assert math.isclose(scaled, raw * temperature**2, abs_tol=1e-12)
        results.append({"temperature": temperature, "teacher_probabilities": p,
                        "student_probabilities": q, "raw_kl": raw, "scaled_kl": scaled,
                        "raw_gradient_norm": norm(raw_grad), "scaled_gradient_norm": norm(grad)})
        print(f"T={temperature:>4.1f} raw_KL={raw:.6f} scaled_KL={scaled:.6f} "
              f"raw_grad={norm(raw_grad):.6f} scaled_grad={norm(grad):.6f}")
    assert max(errors) < 1e-7
    shifted = kl_and_gradient([x+100 for x in teacher], [x-75 for x in student], 4, True)[0]
    assert math.isclose(shifted, results[2]["scaled_kl"], abs_tol=1e-12)
    assert kl_and_gradient(teacher, teacher, 4, True)[0] == 0.0
    forward = kl_and_gradient(teacher, student, 1, False)[0]
    reverse = kl_and_gradient(student, teacher, 1, False)[0]
    assert not math.isclose(forward, reverse)
    teacher_argmax = max(range(3), key=teacher.__getitem__)
    # A local derivative on three free logits, not an optimizer or trained student.
    wrong_teacher = [-2.0, 4.0, 1.0]
    _, wrong_grad, _, _ = kl_and_gradient(wrong_teacher, student, 2, True)
    assert wrong_grad[1] < 0  # gradient descent increases the teacher-favored logit
    print(f"KL(teacher||student)={forward:.6f}; KL(student||teacher)={reverse:.6f}")
    print(f"teacher_argmax={teacher_argmax}; max_gradient_error={max(errors):.3e}")
    print("checks=25 passed; fixed-logit arithmetic only")
    Path("report.json").write_text(json.dumps({"results": results, "gradient_errors": errors,
        "forward_kl": forward, "reverse_kl": reverse, "checks": 25}, indent=2), encoding="utf-8")


if __name__ == "__main__":
    main()

"""Small CPU autograd cross-check against the frozen scalar distillation lesson."""
import json
import math
import platform
import sys
import warnings
from pathlib import Path
import torch
from torch.nn import functional as F

torch.set_num_threads(1)
DTYPE = torch.float64


def main():
    teacher = torch.tensor([[4.0,1.0,-2.0]],dtype=DTYPE,device='cpu')
    rows = []
    print(f'Python {platform.python_version()}; torch {torch.__version__}; device=cpu; dtype=float64; threads=1')
    for temperature in [1.0,2.0,4.0,16.0]:
        student = torch.tensor([[1.0,0.0,-1.0]],dtype=DTYPE,device='cpu',requires_grad=True)
        target = F.softmax(teacher/temperature,dim=-1)
        student_log = F.log_softmax(student/temperature,dim=-1)
        q = student_log.exp()
        raw_loss = F.kl_div(student_log,target,reduction='batchmean')
        scaled_loss = raw_loss * temperature**2
        raw_gradient, = torch.autograd.grad(raw_loss,student,retain_graph=True)
        scaled_gradient, = torch.autograd.grad(scaled_loss,student)
        expected_raw = (q-target)/temperature
        expected_scaled = temperature*(q-target)
        raw_error = (raw_gradient-expected_raw).abs().max().item()
        scaled_error = (scaled_gradient-expected_scaled).abs().max().item()
        assert raw_error < 1e-12 and scaled_error < 1e-12
        assert not teacher.requires_grad and teacher.grad is None
        raw_norm = torch.linalg.vector_norm(raw_gradient).item()
        scaled_norm = torch.linalg.vector_norm(scaled_gradient).item()
        assert math.isclose(scaled_norm,raw_norm*temperature**2,rel_tol=1e-12)
        row = {'temperature':temperature,'kl_batchmean':raw_loss.item(),
            'raw_gradient':raw_gradient.detach().tolist(),'scaled_gradient':scaled_gradient.detach().tolist(),
            'raw_norm':raw_norm,'scaled_norm':scaled_norm,'raw_formula_max_error':raw_error,
            'scaled_formula_max_error':scaled_error}
        rows.append(row)
        print(f'T={temperature:g}: KL={raw_loss.item():.6f}; raw_grad_norm={raw_norm:.6f}; '
              f'scaled_grad_norm={scaled_norm:.6f}; formula_max_error={max(raw_error,scaled_error):.3e}')
    target = F.softmax(teacher/4,dim=-1)
    log_input = F.log_softmax(torch.tensor([[1.0,0.0,-1.0]],dtype=DTYPE)/4,dim=-1)
    batchmean = F.kl_div(log_input,target,reduction='batchmean')
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter('always')
        mean = F.kl_div(log_input,target,reduction='mean')
    assert torch.allclose(mean*3,batchmean,rtol=1e-12,atol=1e-12)
    print(f'reduction at B=1,C=3: batchmean/mean={(batchmean/mean).item():.6f}')
    print('checks passed: 4 raw/scaled autograd comparisons, frozen teacher, reduction factor')
    report = {'python':sys.version,'executable':sys.executable,'torch':torch.__version__,
        'torch_path':torch.__file__,'device':'cpu','dtype':'float64','threads':torch.get_num_threads(),
        'rows':rows,'reduction':{'batchmean':batchmean.item(),'mean':mean.item(),
            'captured_warnings':[str(w.message) for w in caught]},
        'scope':'Actual CPU PyTorch autograd on one fixed row of logits; no network training, optimizer steps, task data or model quality claim.'}
    Path('distillation-torch-report.json').write_text(json.dumps(report,indent=2,allow_nan=False)+'\n',encoding='utf-8')


if __name__ == '__main__':
    main()