Knowledge Distillation Loss

1
#include <iostream>
2
#include <vector>
3
#include <cmath>
4
#include <algorithm>
5
#include <numeric>
6
#include <cassert>
7
8
// Numerically stable log-sum-exp
9
double log_sum_exp(const std::vector<double>& z_over_tau) {
10
    double m = *std::max_element(z_over_tau.begin(), z_over_tau.end());
11
    double sum = 0.0;
12
    for (double v : z_over_tau) sum += std::exp(v - m);
13
    return m + std::log(sum);
14
}
15
16
// Compute softmax(z / tau) in a numerically stable way
17
std::vector<double> softmax_temp(const std::vector<double>& z, double tau) {
18
    std::vector<double> z_over_tau(z.size());
19
    for (size_t i = 0; i < z.size(); ++i) z_over_tau[i] = z[i] / tau;
20
    double lse = log_sum_exp(z_over_tau);
21
    std::vector<double> p(z.size());
22
    for (size_t i = 0; i < z.size(); ++i) p[i] = std::exp(z_over_tau[i] - lse);
23
    return p; // sums to 1
24
}
25
26
// Cross-entropy with hard label y using student probabilities at tau=1
27
// CE = -log p_y; we compute via log-softmax for stability
28
double cross_entropy_hard(const std::vector<double>& logits, int y) {
29
    std::vector<double> z1_over_tau(logits.size());
30
    for (size_t i = 0; i < logits.size(); ++i) z1_over_tau[i] = logits[i]; // tau=1
31
    double lse = log_sum_exp(z1_over_tau);
32
    double log_py = logits[y] - lse; // log-softmax_y
33
    return -log_py;
34
}
35
36
// KL divergence KL(q || p) where q and p are probability vectors
37
// Assumes q_i >= 0, p_i > 0, sum to 1. Adds small epsilon for safety.
38
double kl_divergence(const std::vector<double>& q, const std::vector<double>& p) {
39
    const double eps = 1e-12;
40
    assert(q.size() == p.size());
41
    double kl = 0.0;
42
    for (size_t i = 0; i < q.size(); ++i) {
43
        double qi = std::max(q[i], 0.0);
44
        double pi = std::max(p[i], eps);
45
        if (qi > 0.0) kl += qi * (std::log(qi + eps) - std::log(pi));
46
    }
47
    return kl;
48
}
49
50
// Knowledge Distillation loss for a single example
51
// L = alpha * CE_hard + (1-alpha) * tau^2 * KL(teacher_tau || student_tau)
52
double kd_loss_single(const std::vector<double>& student_logits,
53
                      const std::vector<double>& teacher_logits,
54
                      int y,
55
                      double alpha, double tau) {
56
    double L_hard = cross_entropy_hard(student_logits, y);
57
    std::vector<double> p_s_tau = softmax_temp(student_logits, tau);
58
    std::vector<double> p_t_tau = softmax_temp(teacher_logits, tau);
59
    double L_soft = kl_divergence(p_t_tau, p_s_tau);
60
    return alpha * L_hard + (1.0 - alpha) * (tau * tau) * L_soft;
61
}
62
63
// Batch KD loss: average over batch
64
double kd_loss_batch(const std::vector<std::vector<double>>& S_logits,
65
                     const std::vector<std::vector<double>>& T_logits,
66
                     const std::vector<int>& y,
67
                     double alpha, double tau) {
68
    size_t B = S_logits.size();
69
    double sumL = 0.0;
70
    for (size_t b = 0; b < B; ++b) {
71
        sumL += kd_loss_single(S_logits[b], T_logits[b], y[b], alpha, tau);
72
    }
73
    return sumL / static_cast<double>(B);
74
}
75
76
int main() {
77
    // Example 1: single sample, 3 classes
78
    std::vector<double> student_logits = {2.0, 0.5, -1.0};
79
    std::vector<double> teacher_logits = {3.0, 1.0, -0.5};
80
    int y = 0; // ground-truth class
81
    double alpha = 0.5;
82
    double tau = 3.0;
83
84
    double L_single = kd_loss_single(student_logits, teacher_logits, y, alpha, tau);
85
    std::cout << "Single-sample KD loss = " << L_single << "\n";
86
87
    // Example 2: batch of 2
88
    std::vector<std::vector<double>> S_logits = {
89
        {2.0, 0.5, -1.0},
90
        {0.2, -0.1, 1.2}
91
    };
92
    std::vector<std::vector<double>> T_logits = {
93
        {3.0, 1.0, -0.5},
94
        {-0.2, 0.3, 2.0}
95
    };
96
    std::vector<int> labels = {0, 2};
97
98
    double L_batch = kd_loss_batch(S_logits, T_logits, labels, alpha, tau);
99
    std::cout << "Batch KD loss = " << L_batch << "\n";
100
    return 0;
101
}
102

1
#include <iostream>
2
#include <vector>
3
#include <random>
4
#include <cmath>
5
#include <numeric>
6
#include <algorithm>
7
#include <cassert>
8
9
// Utility: stable log-sum-exp for a vector
10
double log_sum_exp(const std::vector<double>& v) {
11
    double m = *std::max_element(v.begin(), v.end());
12
    double sum = 0.0; for (double x : v) sum += std::exp(x - m);
13
    return m + std::log(sum);
14
}
15
16
std::vector<double> softmax_tau(const std::vector<double>& z, double tau) {
17
    std::vector<double> zt(z.size());
18
    for (size_t i = 0; i < z.size(); ++i) zt[i] = z[i] / tau;
19
    double lse = log_sum_exp(zt);
20
    std::vector<double> p(z.size());
21
    for (size_t i = 0; i < z.size(); ++i) p[i] = std::exp(zt[i] - lse);
22
    return p;
23
}
24
25
// Compute logits: z = x * W, where x is (D), W is (D x C), result is (C)
26
std::vector<double> logits_row(const std::vector<double>& x, const std::vector<std::vector<double>>& W) {
27
    size_t D = x.size(), C = W[0].size();
28
    std::vector<double> z(C, 0.0);
29
    for (size_t c = 0; c < C; ++c) {
30
        double s = 0.0;
31
        for (size_t d = 0; d < D; ++d) s += x[d] * W[d][c];
32
        z[c] = s;
33
    }
34
    return z;
35
}
36
37
// Build one-hot vector of size C for label y
38
std::vector<double> one_hot(size_t C, int y) {
39
    std::vector<double> v(C, 0.0); v[y] = 1.0; return v;
40
}
41
42
// Compute KD loss (batch-averaged) and gradient w.r.t. student logits
43
// grad_z = alpha*(p1 - y_onehot) + (1-alpha)*tau*(p_tau - t_tau)
44
struct KDResult { double loss; std::vector<std::vector<double>> grad_logits; };
45
46
KDResult kd_loss_and_grad(const std::vector<std::vector<double>>& Zs, // B x C student logits
47
                          const std::vector<std::vector<double>>& Zt, // B x C teacher logits
48
                          const std::vector<int>& y, double alpha, double tau) {
49
    size_t B = Zs.size(), C = Zs[0].size();
50
    double loss_sum = 0.0;
51
    std::vector<std::vector<double>> grad(B, std::vector<double>(C, 0.0));
52
53
    for (size_t b = 0; b < B; ++b) {
54
        // Hard CE via log-softmax
55
        double lse1 = log_sum_exp(Zs[b]);
56
        double log_py = Zs[b][y[b]] - lse1;
57
        double Lhard = -log_py;
58
59
        // Probabilities
60
        std::vector<double> p1(C), ptau_s(C), ptau_t(C), oh = one_hot(C, y[b]);
61
        for (size_t c = 0; c < C; ++c) p1[c] = std::exp(Zs[b][c] - lse1); // softmax tau=1
62
        ptau_s = softmax_tau(Zs[b], tau);
63
        ptau_t = softmax_tau(Zt[b], tau);
64
65
        // KL(teacher||student) at tau
66
        double Lsoft = 0.0;
67
        const double eps = 1e-12;
68
        for (size_t c = 0; c < C; ++c) {
69
            double q = std::max(ptau_t[c], 0.0);
70
            double p = std::max(ptau_s[c], eps);
71
            if (q > 0.0) Lsoft += q * (std::log(q + eps) - std::log(p));
72
        }
73
74
        // Total per-sample loss
75
        double L = alpha * Lhard + (1.0 - alpha) * (tau * tau) * Lsoft;
76
        loss_sum += L;
77
78
        // Gradient w.r.t. student logits
79
        for (size_t c = 0; c < C; ++c) {
80
            double g_hard = p1[c] - oh[c];
81
            double g_soft = (ptau_s[c] - ptau_t[c]) * tau; // from tau^2 * KL
82
            grad[b][c] = alpha * g_hard + (1.0 - alpha) * g_soft;
83
        }
84
    }
85
86
    // Average over batch
87
    for (size_t b = 0; b < B; ++b)
88
        for (size_t c = 0; c < C; ++c)
89
            grad[b][c] /= static_cast<double>(B);
90
91
    KDResult res; res.loss = loss_sum / static_cast<double>(B); res.grad_logits = std::move(grad); return res;
92
}
93
94
int main() {
95
    // Synthetic data: D=4 features, C=3 classes
96
    size_t D = 4, C = 3, N = 200; // samples
97
    std::mt19937 rng(42);
98
    std::normal_distribution<double> nd(0.0, 1.0);
99
100
    // Generate features X (N x D) and labels y by a hidden linear rule
101
    std::vector<std::vector<double>> X(N, std::vector<double>(D));
102
    std::vector<int> y(N);
103
    std::vector<std::vector<double>> W_hidden(D, std::vector<double>(C));
104
    for (size_t d = 0; d < D; ++d)
105
        for (size_t c = 0; c < C; ++c)
106
            W_hidden[d][c] = nd(rng);
107
    for (size_t i = 0; i < N; ++i) {
108
        for (size_t d = 0; d < D; ++d) X[i][d] = nd(rng);
109
        // Label from argmax of X * W_hidden (noisy linear separability)
110
        std::vector<double> z = logits_row(X[i], W_hidden);
111
        y[i] = int(std::max_element(z.begin(), z.end()) - z.begin());
112
    }
113
114
    // Teacher: linear model with weights W_t close to hidden rule
115
    std::vector<std::vector<double>> W_t = W_hidden; // idealized teacher
116
117
    // Student: linear model with smaller random weights
118
    std::vector<std::vector<double>> W_s(D, std::vector<double>(C));
119
    std::uniform_real_distribution<double> ur(-0.1, 0.1);
120
    for (size_t d = 0; d < D; ++d)
121
        for (size_t c = 0; c < C; ++c)
122
            W_s[d][c] = ur(rng);
123
124
    // Hyperparameters
125
    double alpha = 0.5, tau = 4.0, lr = 0.1;
126
    size_t epochs = 50, B = 20; // batch size
127
128
    for (size_t epoch = 0; epoch < epochs; ++epoch) {
129
        // Simple SGD over mini-batches
130
        double running_loss = 0.0; size_t batches = 0;
131
        for (size_t start = 0; start < N; start += B) {
132
            size_t end = std::min(N, start + B);
133
            size_t bsz = end - start;
134
            std::vector<std::vector<double>> Zs(bsz, std::vector<double>(C));
135
            std::vector<std::vector<double>> Zt(bsz, std::vector<double>(C));
136
137
            for (size_t i = 0; i < bsz; ++i) {
138
                Zs[i] = logits_row(X[start + i], W_s);
139
                Zt[i] = logits_row(X[start + i], W_t); // teacher is fixed
140
            }
141
142
            // KD loss and grad w.r.t. student logits
143
            std::vector<int> y_batch(y.begin() + start, y.begin() + end);
144
            KDResult res = kd_loss_and_grad(Zs, Zt, y_batch, alpha, tau);
145
            running_loss += res.loss; ++batches;
146
147
            // Backprop to weights: grad_W = X^T * grad_logits / B
148
            std::vector<std::vector<double>> gradW(D, std::vector<double>(C, 0.0));
149
            for (size_t i = 0; i < bsz; ++i) {
150
                for (size_t d = 0; d < D; ++d) {
151
                    for (size_t c = 0; c < C; ++c) {
152
                        gradW[d][c] += X[start + i][d] * res.grad_logits[i][c];
153
                    }
154
                }
155
            }
156
            // SGD update
157
            for (size_t d = 0; d < D; ++d)
158
                for (size_t c = 0; c < C; ++c)
159
                    W_s[d][c] -= lr * gradW[d][c];
160
        }
161
        std::cout << "Epoch " << epoch+1 << ": avg KD loss = " << (running_loss / batches) << "\n";
162
    }
163
164
    // Quick evaluation: accuracy with student
165
    size_t correct = 0;
166
    for (size_t i = 0; i < N; ++i) {
167
        std::vector<double> z = logits_row(X[i], W_s);
168
        int pred = int(std::max_element(z.begin(), z.end()) - z.begin());
169
        if (pred == y[i]) ++correct;
170
    }
171
    std::cout << "Training-set accuracy (student) = " << (100.0 * correct / N) << "%\n";
172
    return 0;
173
}
174