teacher_student_rocket.py

1"""2知识蒸馏粗排（第11章 §11.3）34本文件实现两种粗排蒸馏范式：51. Teacher-Student：精排模型（教师）→ 粗排双塔（学生），KL 散度对齐 logit62. Rocket Launching（火箭发射）：共享 Booster 层，轻量 Light Net 作推理路径78PyTorch autograd 实现，无手写梯度。9"""1011import math12import random13import collections14from typing import List, Tuple, Dict1516import torch17import torch.nn as nn18import torch.nn.functional as F19import torch.optim as optim202122# ─── 通用 MLP ────────────────────────────────────────────────────────────────2324class SimpleMLP(nn.Module):25    def __init__(self, dims: List[int]):26        super().__init__()27        layers = []28        for in_d, out_d in zip(dims[:-1], dims[1:]):29            layers.append(nn.Linear(in_d, out_d))30        self.layers = nn.ModuleList(layers)31        self.output_dim = dims[-1]3233    def forward(self, x: torch.Tensor) -> torch.Tensor:34        h = x35        for i, layer in enumerate(self.layers):36            h = layer(h)37            if i < len(self.layers) - 1:38                h = F.relu(h)39        return h404142# ─── 1. Teacher-Student 蒸馏 ─────────────────────────────────────────────────4344class TeacherModel(nn.Module):45    """大模型（精排教师）：深层 MLP + 交叉特征。"""4647    def __init__(self, n_users: int, n_items: int, emb_dim: int = 32, output_dim: int = 32):48        super().__init__()49        scale = math.sqrt(1.0 / emb_dim)50        self.user_emb = nn.Embedding(n_users, emb_dim)51        self.item_emb = nn.Embedding(n_items, emb_dim)52        self.net = SimpleMLP([emb_dim * 2, 64, 32, output_dim])53        nn.init.uniform_(self.user_emb.weight, -scale, scale)54        nn.init.uniform_(self.item_emb.weight, -scale, scale)5556    def forward(self, uid: torch.Tensor, iid: torch.Tensor) -> torch.Tensor:57        x = torch.cat([self.user_emb(uid), self.item_emb(iid)], dim=-1)58        out = self.net(x)59        return out.norm(dim=-1)  # 模拟打分6061    def soft_labels(self, uid: int, candidates: List[int], temperature: float = 1.0) -> torch.Tensor:62        """对候选集输出 soft label（用于学生蒸馏）。"""63        with torch.no_grad():64            u = torch.tensor([uid]).expand(len(candidates))65            c = torch.tensor(candidates)66            scores = self(u, c) / temperature67            return F.softmax(scores, dim=0)686970class StudentDSSM(nn.Module):71    """轻量双塔学生模型：蒸馏教师的 soft label。"""7273    def __init__(74        self,75        n_users: int,76        n_items: int,77        emb_dim: int = 32,78        output_dim: int = 16,79        temperature: float = 1.0,80    ):81        super().__init__()82        self.n_items = n_items83        self.temperature = temperature84        scale = math.sqrt(1.0 / emb_dim)85        self.user_emb = nn.Embedding(n_users, emb_dim)86        self.item_emb = nn.Embedding(n_items, emb_dim)87        self.user_mlp = SimpleMLP([emb_dim, 32, output_dim])88        self.item_mlp = SimpleMLP([emb_dim, 32, output_dim])89        nn.init.uniform_(self.user_emb.weight, -scale, scale)90        nn.init.uniform_(self.item_emb.weight, -scale, scale)9192    def encode_user(self, uid: torch.Tensor) -> torch.Tensor:93        return self.user_mlp(self.user_emb(uid))9495    def encode_item(self, iid: torch.Tensor) -> torch.Tensor:96        return self.item_mlp(self.item_emb(iid))9798    def score(self, uid: int, iid: int) -> float:99        with torch.no_grad():100            ue = self.encode_user(torch.tensor([uid]))101            ie = self.encode_item(torch.tensor([iid]))102            return F.cosine_similarity(ue, ie).item()103104    def distill_step(105        self,106        uid: int,107        candidates: List[int],108        teacher_probs: torch.Tensor,109        alpha: float = 0.5,110    ) -> torch.Tensor:111        """112        KL 蒸馏 + BPR 硬标签联合训练。113        loss = (1-α)*bpr_loss + α*F.kl_div(log_p_S, p_T, reduction='batchmean')114        """115        u = torch.tensor([uid]).expand(len(candidates))116        c = torch.tensor(candidates)117        ue = self.encode_user(u)                          # [N, D]118        ie = self.encode_item(c)                          # [N, D]119        sims = F.cosine_similarity(ue, ie, dim=1)         # [N]120121        # KL 散度蒸馏122        log_p_S = F.log_softmax(sims / self.temperature, dim=0).unsqueeze(0)123        p_T = teacher_probs.unsqueeze(0)124        kl_loss = F.kl_div(log_p_S, p_T, reduction='batchmean')125126        # BPR 硬标签：取教师打分最高的为正样本127        pos_idx = teacher_probs.argmax().item()128        neg_idx = random.choice([j for j in range(len(candidates)) if j != pos_idx])129        bpr_loss = -F.logsigmoid(sims[pos_idx] - sims[neg_idx])130131        return (1 - alpha) * bpr_loss + alpha * kl_loss132133    def train_distill(134        self,135        user_seqs: Dict[int, List[int]],136        teacher: TeacherModel,137        epochs: int = 6,138        lr: float = 0.01,139    ) -> None:140        optimizer = optim.Adam(self.parameters(), lr=lr)141        for epoch in range(epochs):142            total_loss = 0.0143            for uid, seq in user_seqs.items():144                if not seq:145                    continue146                seen = set(seq)147                pool = [i for i in range(self.n_items) if i not in seen]148                if len(pool) < 2:149                    continue150                candidates = random.sample(pool, k=min(8, len(pool)))151                teacher_probs = teacher.soft_labels(uid, candidates)152153                optimizer.zero_grad()154                loss = self.distill_step(uid, candidates, teacher_probs)155                loss.backward()156                optimizer.step()157                total_loss += loss.item()158            if epoch % 2 == 0:159                print(f"  [Student] epoch {epoch:3d}  avg_loss={total_loss / max(len(user_seqs), 1):.4f}")160161    def pre_rank(162        self, uid: int, candidates: List[int], top_k: int = 10163    ) -> List[Tuple[int, float]]:164        with torch.no_grad():165            u = torch.tensor([uid]).expand(len(candidates))166            c = torch.tensor(candidates)167            ue = self.encode_user(u)168            ie = self.encode_item(c)169            sims = F.cosine_similarity(ue, ie, dim=1)170            scores = [(candidates[j], sims[j].item()) for j in range(len(candidates))]171        scores.sort(key=lambda x: x[1], reverse=True)172        return scores[:top_k]173174175# ─── 2. Rocket Launching（共享 Booster 层）────────────────────────────────────176177class RocketLaunching(nn.Module):178    """179    Rocket Launching 蒸馏范式：180    - Heavy Net（精排路径）= Booster 层 + Heavy Head181    - Light Net（粗排路径）= Booster 层（共享） + Light Head（轻量头）182    训练时让 Light Head 做 hint loss 对齐 Heavy Head 中间层（.detach() 实现 stop-gradient）。183    推理只走 Light Net（Booster + Light Head）。184    """185186    def __init__(187        self,188        n_users: int,189        n_items: int,190        emb_dim: int = 32,191        booster_dim: int = 32,192        output_dim: int = 16,193    ):194        super().__init__()195        self.n_items = n_items196        scale = math.sqrt(1.0 / emb_dim)197        self.user_emb = nn.Embedding(n_users, emb_dim)198        self.item_emb = nn.Embedding(n_items, emb_dim)199        # 共享 Booster200        self.booster = SimpleMLP([emb_dim, booster_dim])201        # Heavy Net Head202        self.heavy_head = SimpleMLP([booster_dim, 64, output_dim])203        # Light Net Head（轻量：单隐层，仍远小于 Heavy Head）204        self.light_head = SimpleMLP([booster_dim, booster_dim, output_dim])205        nn.init.uniform_(self.user_emb.weight, -scale, scale)206        nn.init.uniform_(self.item_emb.weight, -scale, scale)207208    def _booster_out(self, x: torch.Tensor) -> torch.Tensor:209        return self.booster(x)210211    def heavy_encode(self, uid: torch.Tensor, iid: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:212        ub = self._booster_out(self.user_emb(uid))213        ib = self._booster_out(self.item_emb(iid))214        return self.heavy_head(ub), self.heavy_head(ib)215216    def light_encode(self, uid: torch.Tensor, iid: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:217        ub = self._booster_out(self.user_emb(uid))218        ib = self._booster_out(self.item_emb(iid))219        return self.light_head(ub), self.light_head(ib)220221    def forward(222        self,223        uid: torch.Tensor,224        pos_iid: torch.Tensor,225        neg_iid: torch.Tensor,226    ) -> torch.Tensor:227        # Heavy Net BPR loss228        ue_h, pe_h = self.heavy_encode(uid, pos_iid)229        ue_h_neg, ne_h = self.heavy_encode(uid, neg_iid)230        heavy_loss = -F.logsigmoid(231            F.cosine_similarity(ue_h, pe_h, dim=1) - F.cosine_similarity(ue_h_neg, ne_h, dim=1)232        ).mean()233234        # Light Net hint loss：Light head 对齐 Heavy head 输出（L2 归一化后做 MSE）235        # 归一化把向量投到单位球，MSE ∈ [0, 4]，与 BPR loss 量级匹配236        # .detach() 实现 stop-gradient（hint target 不回传梯度给 Heavy head）237        ub = self._booster_out(self.user_emb(uid))238        ue_l = F.normalize(self.light_head(ub), dim=1)239        hint_loss = F.mse_loss(ue_l, F.normalize(ue_h.detach(), dim=1))240241        return heavy_loss + 0.5 * hint_loss242243    def fit(244        self,245        user_seqs: Dict[int, List[int]],246        epochs: int = 6,247        lr: float = 0.01,248    ) -> None:249        """训练入口（命名为 fit 避免遮蔽 nn.Module.train()）。"""250        optimizer = optim.Adam(self.parameters(), lr=lr)251        user_pos: Dict[int, set] = collections.defaultdict(set)252        for uid, seq in user_seqs.items():253            for i in seq:254                user_pos[uid].add(i)255        samples = [256            (uid, seq[t - 1], seq[t])257            for uid, seq in user_seqs.items()258            if len(seq) >= 2259            for t in range(1, len(seq))260        ]261        for epoch in range(epochs):262            random.shuffle(samples)263            total_loss = 0.0264            for uid_s, pos_id, _ in samples:265                neg_id = random.randint(0, self.n_items - 1)266                while neg_id in user_pos[uid_s]:267                    neg_id = random.randint(0, self.n_items - 1)268                uid_t = torch.tensor([uid_s])269                pos_t = torch.tensor([pos_id])270                neg_t = torch.tensor([neg_id])271                optimizer.zero_grad()272                loss = self(uid_t, pos_t, neg_t)273                loss.backward()274                nn.utils.clip_grad_norm_(self.parameters(), max_norm=1.0)275                optimizer.step()276                total_loss += loss.item()277            if epoch % 2 == 0:278                print(f"  [Rocket] epoch {epoch:3d}  avg_loss={total_loss / max(len(samples), 1):.4f}")279280    def light_pre_rank(281        self, uid: int, candidates: List[int], top_k: int = 10282    ) -> List[Tuple[int, float]]:283        """推理只走 Light Net（Booster + Light Head）。"""284        with torch.no_grad():285            ub = self._booster_out(self.user_emb(torch.tensor([uid])))286            ue = self.light_head(ub)  # [1, D]287            cand_t = torch.tensor(candidates)288            ib = self._booster_out(self.item_emb(cand_t))289            ie = self.light_head(ib)  # [N, D]290            sims = F.cosine_similarity(ue, ie, dim=1)291            scores = [(candidates[j], sims[j].item()) for j in range(len(candidates))]292        scores.sort(key=lambda x: x[1], reverse=True)293        return scores[:top_k]294295296# ─── 快速验证 ────────────────────────────────────────────────────────────────297298if __name__ == "__main__":299    random.seed(42)300    torch.manual_seed(42)301302    n_users, n_items = 8, 20303    user_seqs = {u: random.sample(range(n_items), k=random.randint(5, 10)) for u in range(n_users)}304305    print("=== Teacher-Student 蒸馏 ===")306    teacher = TeacherModel(n_users, n_items, emb_dim=16, output_dim=16)307    student = StudentDSSM(n_users, n_items, emb_dim=16, output_dim=16, temperature=1.0)308    student.train_distill(user_seqs, teacher, epochs=6, lr=0.05)309    uid, seq = 0, user_seqs[0]310    cands = [i for i in range(n_items) if i not in set(seq)]311    for iid, s in student.pre_rank(uid, cands, top_k=5):312        print(f"  item {iid}: {s:.4f}")313314    print("\n=== Rocket Launching ===")315    rocket = RocketLaunching(n_users, n_items, emb_dim=16, booster_dim=16, output_dim=16)316    rocket.fit(user_seqs, epochs=6, lr=0.05)317    for iid, s in rocket.light_pre_rank(uid, cands, top_k=5):318        print(f"  item {iid}: {s:.4f}")