推荐系统（五）—— Embedding表示学习

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from collections import Counter
import random


class Item2Vec(nn.Module):
    """Skip-gram + Negative Sampling

    两套 Embedding：
      - input_embeddings：推理时使用，是物品的"真实"向量
      - output_embeddings：训练辅助，训练完丢弃
    """

    def __init__(self, vocab_size: int, embedding_dim: int, num_negatives: int = 5):
        super().__init__()
        self.input_embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.output_embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.num_negatives = num_negatives

        # Xavier 初始化让初始内积处于较小范围，
        # 避免第一轮迭代时 sigmoid 已经饱和。
        nn.init.xavier_uniform_(self.input_embeddings.weight)
        nn.init.xavier_uniform_(self.output_embeddings.weight)

    def forward(self, target, context, negatives):
        """
        target    : (B,)          中心物品 id
        context   : (B,)          真实上下文 id
        negatives : (B, K)        每行 K 个随机负样本
        """
        t = self.input_embeddings(target)      # (B, d)
        c = self.output_embeddings(context)    # (B, d)
        n = self.output_embeddings(negatives)  # (B, K, d)

        # 正样本：把中心和真实上下文的内积推大
        pos = torch.clamp((t * c).sum(-1), -10, 10)
        pos_loss = -torch.log(torch.sigmoid(pos) + 1e-10)

        # 负样本：把中心和 K 个随机物品的内积推小
        neg = torch.clamp(torch.bmm(n, t.unsqueeze(-1)).squeeze(-1), -10, 10)
        neg_loss = -torch.log(torch.sigmoid(-neg) + 1e-10).sum(-1)

        return (pos_loss + neg_loss).mean()

    @torch.no_grad()
    def vector(self, item_id: int) -> np.ndarray:
        return self.input_embeddings.weight[item_id].cpu().numpy()


# --- 数据流水线 ----------------------------------------------------------- #

def build_vocab(sequences):
    counter = Counter(item for seq in sequences for item in seq)
    vocab = {item: idx for idx, item in enumerate(counter)}
    return vocab, counter


def make_pairs(sequences, vocab, window=5):
    pairs = []
    for seq in sequences:
        ids = [vocab[i] for i in seq if i in vocab]
        for t, centre in enumerate(ids):
            lo, hi = max(0, t - window), min(len(ids), t + window + 1)
            for j in range(lo, hi):
                if j != t:
                    pairs.append((centre, ids[j]))
    return pairs


def make_neg_sampler(counter, vocab):
    items = list(counter)
    probs = np.array([counter[i] ** 0.75 for i in items], dtype=np.float64)
    probs /= probs.sum()
    idx_lookup = np.array([vocab[i] for i in items])

    def sample(n_rows: int, k: int) -> np.ndarray:
        chosen = np.random.choice(len(items), size=(n_rows, k), p=probs)
        return idx_lookup[chosen]

    return sample


# --- 训练循环 ------------------------------------------------------------- #

if __name__ == "__main__":
    sequences = [
        [1, 2, 3, 4, 5], [2, 3, 4, 6, 7],
        [1, 3, 5, 8, 9], [4, 5, 6, 10, 11],
    ]
    vocab, counter = build_vocab(sequences)
    pairs = make_pairs(sequences, vocab, window=2)
    neg = make_neg_sampler(counter, vocab)

    model = Item2Vec(len(vocab), embedding_dim=64, num_negatives=5)
    opt = optim.Adam(model.parameters(), lr=1e-3)

    BATCH = 32
    for epoch in range(10):
        random.shuffle(pairs)
        running = 0.0
        for i in range(0, len(pairs), BATCH):
            batch = pairs[i:i + BATCH]
            t = torch.tensor([p[0] for p in batch], dtype=torch.long)
            c = torch.tensor([p[1] for p in batch], dtype=torch.long)
            ng = torch.tensor(neg(len(batch), 5), dtype=torch.long)

            opt.zero_grad()
            loss = model(t, c, ng)
            loss.backward()
            opt.step()
            running += loss.item()
        print(f"epoch {epoch + 1:>2}  loss {running / max(1, len(pairs) // BATCH):.4f}")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

class Item2VecCBOW(nn.Module):
    """从平均的上下文窗口预测中心物品"""

    def __init__(self, vocab_size, embedding_dim, num_negatives=5):
        super().__init__()
        self.in_emb = nn.Embedding(vocab_size, embedding_dim)
        self.out_emb = nn.Embedding(vocab_size, embedding_dim)
        nn.init.xavier_uniform_(self.in_emb.weight)
        nn.init.xavier_uniform_(self.out_emb.weight)

    def forward(self, context, target, negatives):
        ctx = self.in_emb(context).mean(dim=1)        # (B, d)
        tgt = self.out_emb(target)                    # (B, d)
        neg = self.out_emb(negatives)                 # (B, K, d)

        pos = torch.clamp((ctx * tgt).sum(-1), -10, 10)
        pos_loss = -torch.log(torch.sigmoid(pos) + 1e-10)

        neg = torch.clamp(torch.bmm(neg, ctx.unsqueeze(-1)).squeeze(-1), -10, 10)
        neg_loss = -torch.log(torch.sigmoid(-neg) + 1e-10).sum(-1)
        return (pos_loss + neg_loss).mean()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

import numpy as np
import networkx as nx
import random


class Node2Vec:
    """有偏随机游走 + Skip-gram，输入是 NetworkX 图"""

    def __init__(self, graph, dimensions=64, walk_length=80,
                 num_walks=10, p=1.0, q=1.0, window_size=10):
        self.graph = graph
        self.dimensions = dimensions
        self.walk_length = walk_length
        self.num_walks = num_walks
        self.p = p
        self.q = q
        self.window_size = window_size

    def _walk(self, start):
        walk = [start]
        while len(walk) < self.walk_length:
            cur = walk[-1]
            nbrs = list(self.graph.neighbors(cur))
            if not nbrs:
                break

            if len(walk) == 1:
                walk.append(random.choice(nbrs))
                continue

            prev = walk[-2]
            weights = []
            for x in nbrs:
                w = self.graph[cur][x].get("weight", 1.0)
                if x == prev:                         # 距离 0
                    weights.append(w / self.p)
                elif self.graph.has_edge(prev, x):    # 距离 1
                    weights.append(w)
                else:                                 # 距离 2
                    weights.append(w / self.q)
            weights = np.array(weights)
            weights /= weights.sum()
            walk.append(np.random.choice(nbrs, p=weights))
        return walk

    def generate_walks(self):
        walks = []
        nodes = list(self.graph.nodes())
        for _ in range(self.num_walks):
            random.shuffle(nodes)
            for n in nodes:
                walks.append(self._walk(n))
        return walks

    def fit(self):
        from gensim.models import Word2Vec
        walks = [[str(n) for n in w] for w in self.generate_walks()]
        model = Word2Vec(walks, vector_size=self.dimensions,
                         window=self.window_size, sg=1,
                         negative=5, min_count=0, workers=4)
        return {n: model.wv[str(n)] for n in self.graph.nodes()}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

from collections import defaultdict

def co_occurrence_graph(interactions, min_jaccard=0.1, max_users=None):
    """从 (user, item) 交互构造物品-物品图。
    边权 = 同时与两件物品发生过交互的用户集合的 Jaccard 相似度。
    """
    item_users = defaultdict(set)
    for u, i in interactions:
        item_users[i].add(u)

    G = nx.Graph()
    items = list(item_users)
    for a in range(len(items)):
        ua = item_users[items[a]]
        if max_users and len(ua) > max_users:
            continue
        G.add_node(items[a])
        for b in range(a + 1, len(items)):
            ub = item_users[items[b]]
            inter = len(ua & ub)
            if inter == 0:
                continue
            j = inter / len(ua | ub)
            if j >= min_jaccard:
                G.add_edge(items[a], items[b], weight=j)
    return G

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

import torch
import torch.nn as nn
import torch.nn.functional as F


def make_tower(in_dim: int, hidden: list[int], out_dim: int) -> nn.Sequential:
    layers, d = [], in_dim
    for h in hidden:
        layers += [nn.Linear(d, h), nn.ReLU(), nn.BatchNorm1d(h)]
        d = h
    layers.append(nn.Linear(d, out_dim))
    return nn.Sequential(*layers)


class DSSM(nn.Module):
    """对称双塔 + 余弦相似度打分"""

    def __init__(self, user_dim, item_dim, embedding_dim=128, hidden=(256, 128)):
        super().__init__()
        self.user_tower = make_tower(user_dim, list(hidden), embedding_dim)
        self.item_tower = make_tower(item_dim, list(hidden), embedding_dim)

    def encode_user(self, x):
        return F.normalize(self.user_tower(x), p=2, dim=-1)

    def encode_item(self, x):
        return F.normalize(self.item_tower(x), p=2, dim=-1)

    def forward(self, user_x, item_x):
        u = self.encode_user(user_x)
        i = self.encode_item(item_x)
        return (u * i).sum(-1)        # 已 L2 归一化，内积即余弦


class SampledSoftmaxLoss(nn.Module):
    """正样本必须在 in-batch 负样本中胜出（带温度的交叉熵）"""

    def __init__(self, temperature: float = 0.1):
        super().__init__()
        self.t = temperature

    def forward(self, u, pos, neg):
        # u   : (B, d)
        # pos : (B, d)
        # neg : (B, K, d)
        pos_score = (u * pos).sum(-1, keepdim=True) / self.t          # (B, 1)
        neg_score = torch.bmm(u.unsqueeze(1),
                              neg.transpose(1, 2)).squeeze(1) / self.t  # (B, K)
        logits = torch.cat([pos_score, neg_score], dim=1)             # (B, 1+K)
        target = torch.zeros(u.size(0), dtype=torch.long, device=u.device)
        return F.cross_entropy(logits, target)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

class YouTubeDNN(nn.Module):
    """对观看历史做平均池化，拼上侧信息特征，再投影到 d 维"""

    def __init__(self, num_videos, num_categories, d=64,
                 user_hidden=(256, 128), item_hidden=(128, 64)):
        super().__init__()
        self.video = nn.Embedding(num_videos, d)
        self.category = nn.Embedding(num_categories, 16)

        self.user_mlp = make_tower(d + 16, list(user_hidden), d)
        self.item_mlp = make_tower(d + 16, list(item_hidden), d)

    def user_vector(self, history_ids, history_cats):
        h = self.video(history_ids).mean(dim=1)
        c = self.category(history_cats).mean(dim=1)
        return F.normalize(self.user_mlp(torch.cat([h, c], dim=-1)), dim=-1)

    def item_vector(self, item_id, item_cat):
        e = self.video(item_id)
        c = self.category(item_cat)
        return F.normalize(self.item_mlp(torch.cat([e, c], dim=-1)), dim=-1)

    def forward(self, hist_ids, hist_cats, item_id, item_cat):
        u = self.user_vector(hist_ids, hist_cats)
        i = self.item_vector(item_id, item_cat)
        return (u * i).sum(-1)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

class YouTubeDNNWithAttention(nn.Module):
    def __init__(self, num_videos, num_categories, d=64, n_heads=4):
        super().__init__()
        self.video = nn.Embedding(num_videos, d)
        self.category = nn.Embedding(num_categories, 16)
        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
        self.user_mlp = make_tower(d + 16, [256], d)
        self.item_mlp = make_tower(d + 16, [128], d)

    def user_vector(self, history_ids, history_cats):
        h = self.video(history_ids)                        # (B, T, d)
        h, _ = self.attn(h, h, h)                          # 自注意力
        h = h.mean(dim=1)                                  # 对加权后的序列再池化
        c = self.category(history_cats).mean(dim=1)
        return F.normalize(self.user_mlp(torch.cat([h, c], dim=-1)), dim=-1)

    def item_vector(self, item_id, item_cat):
        e = self.video(item_id)
        c = self.category(item_cat)
        return F.normalize(self.item_mlp(torch.cat([e, c], dim=-1)), dim=-1)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

import numpy as np

def uniform_neg(catalog_size: int, k: int, exclude: set | None = None) -> list[int]:
    out = set()
    while len(out) < k:
        i = int(np.random.randint(catalog_size))
        if exclude is None or i not in exclude:
            out.add(i)
    return list(out)


def popularity_neg(probs: np.ndarray, k: int) -> np.ndarray:
    """probs 是长度为 N 的、已经按 f^0.75 归一化的概率向量"""
    return np.random.choice(len(probs), size=k, replace=False, p=probs)


def in_batch_neg(item_emb: torch.Tensor) -> torch.Tensor:
    """item_emb: (B, d)。返回 (B, B-1, d)，即用 batch 中其他行作为负样本"""
    B = item_emb.size(0)
    eye = torch.eye(B, dtype=torch.bool, device=item_emb.device)
    idx = (~eye).nonzero(as_tuple=False)[:, 1].view(B, B - 1)
    return item_emb[idx]


def hard_neg(user_emb, item_emb, positives_mask, top_k=100, k=10):
    """Hard negative：与用户最相似但不在正样本里的物品"""
    sims = user_emb @ item_emb.T
    sims[positives_mask] = -1e9
    top = sims.topk(top_k, dim=-1).indices
    pick = torch.randint(0, top_k, (user_emb.size(0), k), device=user_emb.device)
    return top.gather(1, pick)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

import faiss
import numpy as np


class ANNIndex:
    """对 FAISS 三种常用索引的薄封装"""

    def __init__(self, dim: int, kind: str = "IVF",
                 nlist: int = 100, hnsw_m: int = 32):
        if kind == "Flat":
            self.index = faiss.IndexFlatIP(dim)
        elif kind == "IVF":
            quantiser = faiss.IndexFlatIP(dim)
            self.index = faiss.IndexIVFFlat(
                quantiser, dim, nlist, faiss.METRIC_INNER_PRODUCT)
        elif kind == "HNSW":
            self.index = faiss.IndexHNSWFlat(dim, hnsw_m,
                                             faiss.METRIC_INNER_PRODUCT)
        else:
            raise ValueError(kind)
        self._trained = False

    def build(self, vectors: np.ndarray) -> None:
        x = np.ascontiguousarray(vectors, dtype="float32")
        faiss.normalize_L2(x)             # 内积 + 归一化 = 余弦
        if hasattr(self.index, "train") and not self._trained:
            self.index.train(x)
            self._trained = True
        self.index.add(x)

    def query(self, vector: np.ndarray, k: int = 10):
        q = np.ascontiguousarray(vector.reshape(1, -1), dtype="float32")
        faiss.normalize_L2(q)
        sims, ids = self.index.search(q, k)
        return sims[0], ids[0]

    def set_nprobe(self, n: int) -> None:
        """IVF 专用：每次查询扫多少个簇。值越大，越慢但召回越高"""
        if hasattr(self.index, "nprobe"):
            self.index.nprobe = n

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

import numpy as np
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score


def neighbour_overlap(emb: dict, gold: dict, k: int = 10) -> float:
    """Embedding top-k 和 ground-truth top-k 的 Jaccard 平均"""
    scores = []
    items = list(emb)
    matrix = np.stack([emb[i] for i in items])
    matrix /= np.linalg.norm(matrix, axis=1, keepdims=True) + 1e-9

    sims = matrix @ matrix.T
    np.fill_diagonal(sims, -np.inf)

    for idx, item in enumerate(items):
        if item not in gold:
            continue
        top_emb = set(items[j] for j in np.argsort(-sims[idx])[:k])
        top_gold = set(sorted(gold[item], key=gold[item].get, reverse=True)[:k])
        scores.append(len(top_emb & top_gold) / k)
    return float(np.mean(scores)) if scores else 0.0


def cluster_silhouette(emb: dict, n_clusters: int = 10) -> float:
    matrix = np.stack(list(emb.values()))
    labels = KMeans(n_clusters=n_clusters, random_state=0, n_init=10).fit_predict(matrix)
    return silhouette_score(matrix, labels)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

def hit_rate_ndcg(user_vec: dict, item_vec: dict, test, k: int = 10):
    items = list(item_vec)
    matrix = np.stack([item_vec[i] for i in items])
    item_index = {i: idx for idx, i in enumerate(items)}

    hits = 0
    ndcg = []
    for user_id, true_item in test:
        if user_id not in user_vec or true_item not in item_index:
            continue
        sims = matrix @ user_vec[user_id]
        top = np.argsort(-sims)[:k]
        if item_index[true_item] in top:
            hits += 1
            rank = int(np.where(top == item_index[true_item])[0][0]) + 1
            ndcg.append(1.0 / np.log2(rank + 1))
        else:
            ndcg.append(0.0)
    return {"hit_rate@k": hits / len(test), "ndcg@k": float(np.mean(ndcg))}