Skip to content

Distillers

Distillers define how teacher knowledge is transferred to the student model. RecDistillery currently includes embedding-based, ranking-based, and composite distillation strategies.

Available Distillers

DE
RRD
DE_RRD
HTD
FTD
UnKD

Distiller defaults are stored in:

config/distillation/

Distiller Base

Distiller

Bases: Module, ABC

Source code in recdistill/distillers/base.py 
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
class Distiller(nn.Module, ABC):
    def on_train_start(
        self,
        teacher_state: TeacherState,
        dataset: InteractionDataset,
    ) -> None:
        del teacher_state, dataset

    def on_epoch_start(self) -> None:
        return None

    def build_aux_batch(
        self,
        batch: InteractionBatch,
        device: torch.device,
    ) -> object | None:
        del batch, device
        return None

    def compute_loss(
        self,
        student: nn.Module,
        batch: InteractionBatch,
        aux_batch: object | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError

on_train_start(teacher_state: TeacherState, dataset: InteractionDataset) -> None

Source code in recdistill/distillers/base.py 
14
15
16
17
18
19
def on_train_start(
    self,
    teacher_state: TeacherState,
    dataset: InteractionDataset,
) -> None:
    del teacher_state, dataset

on_epoch_start() -> None

Source code in recdistill/distillers/base.py 
21
22
def on_epoch_start(self) -> None:
    return None

build_aux_batch(batch: InteractionBatch, device: torch.device) -> object | None

Source code in recdistill/distillers/base.py 
24
25
26
27
28
29
30
def build_aux_batch(
    self,
    batch: InteractionBatch,
    device: torch.device,
) -> object | None:
    del batch, device
    return None

compute_loss(student: nn.Module, batch: InteractionBatch, aux_batch: object | None = None) -> torch.Tensor

Source code in recdistill/distillers/base.py 
32
33
34
35
36
37
38
def compute_loss(
    self,
    student: nn.Module,
    batch: InteractionBatch,
    aux_batch: object | None = None,
) -> torch.Tensor:
    raise NotImplementedError

DE

Expert

Bases: Module

Source code in recdistill/distillers/de.py 
11
12
13
14
15
16
17
18
19
20
21
class Expert(nn.Module):
    def __init__(self, dims: list[int]):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(dims[0], dims[1]),
            nn.ReLU(),
            nn.Linear(dims[1], dims[2]),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.mlp(inputs)

mlp = nn.Sequential(nn.Linear(dims[0], dims[1]), nn.ReLU(), nn.Linear(dims[1], dims[2])) instance-attribute

__init__(dims: list[int])

Source code in recdistill/distillers/de.py 
12
13
14
15
16
17
18
def __init__(self, dims: list[int]):
    super().__init__()
    self.mlp = nn.Sequential(
        nn.Linear(dims[0], dims[1]),
        nn.ReLU(),
        nn.Linear(dims[1], dims[2]),
    )

forward(inputs: torch.Tensor) -> torch.Tensor

Source code in recdistill/distillers/de.py 
20
21
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
    return self.mlp(inputs)

DEDistiller

Bases: Distiller

Source code in recdistill/distillers/de.py 
 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
116
117
118
class DEDistiller(Distiller):
    def __init__(
        self,
        teacher_dim: int,
        student_dim: int,
        num_experts: int,
        lambda_de: float,
        temperature: float = 0.01,
    ):
        super().__init__()
        if num_experts < 1:
            raise ValueError("num_experts must be >= 1.")
        if temperature <= 0:
            raise ValueError("temperature must be > 0.")
        if lambda_de < 0:
            raise ValueError("lambda_de must be >= 0.")

        self.teacher_dim = teacher_dim
        self.student_dim = student_dim
        self.num_experts = num_experts
        self.lambda_de = lambda_de
        self.temperature = temperature

        hidden_dim = (teacher_dim + student_dim) // 2
        if teacher_dim == student_dim:
            hidden_dim = max(1, student_dim // 2)

        dims = [student_dim, hidden_dim, teacher_dim]
        self.user_experts = nn.ModuleList(Expert(dims) for _ in range(num_experts))
        self.item_experts = nn.ModuleList(Expert(dims) for _ in range(num_experts))
        self.user_gate = nn.Sequential(nn.Linear(teacher_dim, num_experts), nn.Softmax(dim=1))
        self.item_gate = nn.Sequential(nn.Linear(teacher_dim, num_experts), nn.Softmax(dim=1))

        self.register_buffer("_teacher_users", torch.empty(0), persistent=False)
        self.register_buffer("_teacher_items", torch.empty(0), persistent=False)
        self.softmax = nn.Softmax(dim=1)

    def on_train_start(self, teacher_state: TeacherState, dataset) -> None:
        del dataset
        device = self._teacher_users.device if self._teacher_users.numel() else self.user_experts[0].mlp[0].weight.device
        self._teacher_users = teacher_state.user_embeddings.detach().to(device)
        self._teacher_items = teacher_state.item_embeddings.detach().to(device)

    def set_temperature(self, temperature: float) -> None:
        self.temperature = temperature

    def compute_loss(
        self,
        student: nn.Module,
        batch: InteractionBatch,
        aux_batch: object | None = None,
    ) -> torch.Tensor:
        del aux_batch
        user_loss = self._entity_loss(
            student_indices=batch.unique_users,
            student_table=student.get_all_user_embeddings(),
            teacher_table=self._teacher_users,
            experts=self.user_experts,
            gate=self.user_gate,
        )
        item_loss = self._entity_loss(
            student_indices=batch.unique_items,
            student_table=student.get_all_item_embeddings(),
            teacher_table=self._teacher_items,
            experts=self.item_experts,
            gate=self.item_gate,
        )
        return self.lambda_de * (user_loss + item_loss)

    def _entity_loss(
        self,
        student_indices: torch.Tensor,
        student_table: torch.Tensor,
        teacher_table: torch.Tensor,
        experts: nn.ModuleList,
        gate: nn.Module,
    ) -> torch.Tensor:
        student_emb = student_table[student_indices]
        teacher_emb = teacher_table[student_indices]
        selection = gate(teacher_emb)

        if self.num_experts == 1:
            selection_result = 1.0
        else:
            noise = torch.distributions.Gumbel(0, 1).sample(selection.size()).to(selection.device)
            selection = selection + 1e-10
            selection = self.softmax((selection.log() + noise) / self.temperature)
            selection = selection.unsqueeze(1).repeat(1, self.teacher_dim, 1)
            selection_result = selection

        expert_outputs = [experts[idx](student_emb).unsqueeze(-1) for idx in range(self.num_experts)]
        expert_outputs = torch.cat(expert_outputs, dim=-1)
        mixed = expert_outputs * selection_result
        mixed = mixed.sum(dim=2)
        return ((teacher_emb - mixed) ** 2).sum(dim=-1).mean()

teacher_dim = teacher_dim instance-attribute

student_dim = student_dim instance-attribute

num_experts = num_experts instance-attribute

lambda_de = lambda_de instance-attribute

temperature = temperature instance-attribute

user_experts = nn.ModuleList((Expert(dims)) for _ in (range(num_experts))) instance-attribute

item_experts = nn.ModuleList((Expert(dims)) for _ in (range(num_experts))) instance-attribute

user_gate = nn.Sequential(nn.Linear(teacher_dim, num_experts), nn.Softmax(dim=1)) instance-attribute

item_gate = nn.Sequential(nn.Linear(teacher_dim, num_experts), nn.Softmax(dim=1)) instance-attribute

softmax = nn.Softmax(dim=1) instance-attribute

__init__(teacher_dim: int, student_dim: int, num_experts: int, lambda_de: float, temperature: float = 0.01)

Source code in recdistill/distillers/de.py 
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
def __init__(
    self,
    teacher_dim: int,
    student_dim: int,
    num_experts: int,
    lambda_de: float,
    temperature: float = 0.01,
):
    super().__init__()
    if num_experts < 1:
        raise ValueError("num_experts must be >= 1.")
    if temperature <= 0:
        raise ValueError("temperature must be > 0.")
    if lambda_de < 0:
        raise ValueError("lambda_de must be >= 0.")

    self.teacher_dim = teacher_dim
    self.student_dim = student_dim
    self.num_experts = num_experts
    self.lambda_de = lambda_de
    self.temperature = temperature

    hidden_dim = (teacher_dim + student_dim) // 2
    if teacher_dim == student_dim:
        hidden_dim = max(1, student_dim // 2)

    dims = [student_dim, hidden_dim, teacher_dim]
    self.user_experts = nn.ModuleList(Expert(dims) for _ in range(num_experts))
    self.item_experts = nn.ModuleList(Expert(dims) for _ in range(num_experts))
    self.user_gate = nn.Sequential(nn.Linear(teacher_dim, num_experts), nn.Softmax(dim=1))
    self.item_gate = nn.Sequential(nn.Linear(teacher_dim, num_experts), nn.Softmax(dim=1))

    self.register_buffer("_teacher_users", torch.empty(0), persistent=False)
    self.register_buffer("_teacher_items", torch.empty(0), persistent=False)
    self.softmax = nn.Softmax(dim=1)

on_train_start(teacher_state: TeacherState, dataset) -> None

Source code in recdistill/distillers/de.py 
61
62
63
64
65
def on_train_start(self, teacher_state: TeacherState, dataset) -> None:
    del dataset
    device = self._teacher_users.device if self._teacher_users.numel() else self.user_experts[0].mlp[0].weight.device
    self._teacher_users = teacher_state.user_embeddings.detach().to(device)
    self._teacher_items = teacher_state.item_embeddings.detach().to(device)

set_temperature(temperature: float) -> None

Source code in recdistill/distillers/de.py 
67
68
def set_temperature(self, temperature: float) -> None:
    self.temperature = temperature

compute_loss(student: nn.Module, batch: InteractionBatch, aux_batch: object | None = None) -> torch.Tensor

Source code in recdistill/distillers/de.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def compute_loss(
    self,
    student: nn.Module,
    batch: InteractionBatch,
    aux_batch: object | None = None,
) -> torch.Tensor:
    del aux_batch
    user_loss = self._entity_loss(
        student_indices=batch.unique_users,
        student_table=student.get_all_user_embeddings(),
        teacher_table=self._teacher_users,
        experts=self.user_experts,
        gate=self.user_gate,
    )
    item_loss = self._entity_loss(
        student_indices=batch.unique_items,
        student_table=student.get_all_item_embeddings(),
        teacher_table=self._teacher_items,
        experts=self.item_experts,
        gate=self.item_gate,
    )
    return self.lambda_de * (user_loss + item_loss)

RRD

RRDDistiller

Bases: Distiller

Source code in recdistill/distillers/rrd.py 
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
class RRDDistiller(Distiller):
    def __init__(self, sampler: RRDSampler, lambda_rrd: float):
        super().__init__()
        self.sampler = sampler
        self.lambda_rrd = lambda_rrd

    def on_train_start(self, teacher_state, dataset) -> None:
        self.sampler.initialize(dataset=dataset, teacher_state=teacher_state)

    def on_epoch_start(self) -> None:
        self.sampler.refresh()

    def build_aux_batch(
        self,
        batch: InteractionBatch,
        device: torch.device,
    ) -> RRDAuxBatch:
        return self.sampler.sample(batch.unique_users, device=device)

    def compute_loss(
        self,
        student: torch.nn.Module,
        batch: InteractionBatch,
        aux_batch: RRDAuxBatch | None = None,
    ) -> torch.Tensor:
        if aux_batch is None:
            return torch.zeros((), device=batch.users.device)

        if _can_score_rrd_items_together(student):
            items = torch.cat([aux_batch.interesting_items, aux_batch.uninteresting_items], dim=1)
            scores = student.score_items(aux_batch.users, items)
            interesting_scores, uninteresting_scores = scores.split(
                [aux_batch.interesting_items.size(1), aux_batch.uninteresting_items.size(1)],
                dim=1,
            )
        else:
            interesting_scores = student.score_items(aux_batch.users, aux_batch.interesting_items)
            uninteresting_scores = student.score_items(aux_batch.users, aux_batch.uninteresting_items)
        return self.lambda_rrd * relaxed_ranking_loss(interesting_scores, uninteresting_scores)

sampler = sampler instance-attribute

lambda_rrd = lambda_rrd instance-attribute

__init__(sampler: RRDSampler, lambda_rrd: float)

Source code in recdistill/distillers/rrd.py 
30
31
32
33
def __init__(self, sampler: RRDSampler, lambda_rrd: float):
    super().__init__()
    self.sampler = sampler
    self.lambda_rrd = lambda_rrd

on_train_start(teacher_state, dataset) -> None

Source code in recdistill/distillers/rrd.py 
35
36
def on_train_start(self, teacher_state, dataset) -> None:
    self.sampler.initialize(dataset=dataset, teacher_state=teacher_state)

on_epoch_start() -> None

Source code in recdistill/distillers/rrd.py 
38
39
def on_epoch_start(self) -> None:
    self.sampler.refresh()

build_aux_batch(batch: InteractionBatch, device: torch.device) -> RRDAuxBatch

Source code in recdistill/distillers/rrd.py 
41
42
43
44
45
46
def build_aux_batch(
    self,
    batch: InteractionBatch,
    device: torch.device,
) -> RRDAuxBatch:
    return self.sampler.sample(batch.unique_users, device=device)

compute_loss(student: torch.nn.Module, batch: InteractionBatch, aux_batch: RRDAuxBatch | None = None) -> torch.Tensor

Source code in recdistill/distillers/rrd.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def compute_loss(
    self,
    student: torch.nn.Module,
    batch: InteractionBatch,
    aux_batch: RRDAuxBatch | None = None,
) -> torch.Tensor:
    if aux_batch is None:
        return torch.zeros((), device=batch.users.device)

    if _can_score_rrd_items_together(student):
        items = torch.cat([aux_batch.interesting_items, aux_batch.uninteresting_items], dim=1)
        scores = student.score_items(aux_batch.users, items)
        interesting_scores, uninteresting_scores = scores.split(
            [aux_batch.interesting_items.size(1), aux_batch.uninteresting_items.size(1)],
            dim=1,
        )
    else:
        interesting_scores = student.score_items(aux_batch.users, aux_batch.interesting_items)
        uninteresting_scores = student.score_items(aux_batch.users, aux_batch.uninteresting_items)
    return self.lambda_rrd * relaxed_ranking_loss(interesting_scores, uninteresting_scores)

relaxed_ranking_loss(interesting_scores: torch.Tensor, uninteresting_scores: torch.Tensor) -> torch.Tensor

Source code in recdistill/distillers/rrd.py 
10
11
12
13
14
15
def relaxed_ranking_loss(
    interesting_scores: torch.Tensor,
    uninteresting_scores: torch.Tensor,
) -> torch.Tensor:
    diff = interesting_scores.unsqueeze(-1) - uninteresting_scores.unsqueeze(1)
    return -torch.nn.functional.logsigmoid(diff).mean()

UnKD

UnKDDistiller

Bases: Distiller

UnKD distillation from the standalone implementation in UnKD/sample.py.

The distiller samples teacher-ranked item pairs inside popularity groups, then asks the student to preserve the teacher preference ordering with a BPR-style loss. Popularity groups and ratios are built from the recdistill InteractionDataset so the original UnKD loader is not required.

Source code in recdistill/distillers/unkd.py 
 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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
class UnKDDistiller(Distiller):
    """
    UnKD distillation from the standalone implementation in UnKD/sample.py.

    The distiller samples teacher-ranked item pairs inside popularity groups, then
    asks the student to preserve the teacher preference ordering with a BPR-style
    loss. Popularity groups and ratios are built from the recdistill
    InteractionDataset so the original UnKD loader is not required.
    """

    def __init__(
        self,
        lambda_unkd: float = 1.0,
        sample_num: int = 30,
        group_count: int = 2,
        popularity_lambda: float = 1.0,
        rank_top_k: int = 1000,
        rank_temperature: float = 20.0,
    ):
        super().__init__()
        if lambda_unkd < 0:
            raise ValueError("lambda_unkd must be >= 0.")
        if sample_num < 1:
            raise ValueError("sample_num must be >= 1.")
        if group_count < 1:
            raise ValueError("group_count must be >= 1.")
        if popularity_lambda < 0:
            raise ValueError("popularity_lambda must be >= 0.")
        if rank_top_k < 1:
            raise ValueError("rank_top_k must be >= 1.")
        if rank_temperature <= 0:
            raise ValueError("rank_temperature must be > 0.")

        self.lambda_unkd = float(lambda_unkd)
        self.sample_num = int(sample_num)
        self.group_count = int(group_count)
        self.popularity_lambda = float(popularity_lambda)
        self.rank_top_k = int(rank_top_k)
        self.rank_temperature = float(rank_temperature)

        self.dataset: InteractionDataset | None = None
        self.teacher_state: TeacherState | None = None
        self._group_items: list[torch.Tensor] = []
        self._group_sample_counts: list[int] = []
        self._ranked_group_items: list[list[torch.Tensor]] = []
        self._ranked_group_weights: list[list[torch.Tensor]] = []
        self._epoch_pos_items: torch.Tensor | None = None
        self._epoch_neg_items: torch.Tensor | None = None
        self._device_cache: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {}

    def on_train_start(self, teacher_state: TeacherState, dataset: InteractionDataset) -> None:
        self.teacher_state = teacher_state
        self.dataset = dataset
        self._group_items = self._build_popularity_groups(dataset)
        group_ratios = self._build_group_ratios(dataset, self._group_items)
        self._group_sample_counts = self._allocate_group_samples(group_ratios)
        self._ranked_group_items = self._build_ranked_group_items(teacher_state, dataset)
        self._ranked_group_weights = self._build_ranked_group_weights(self._ranked_group_items)
        self.refresh()

    def on_epoch_start(self) -> None:
        self.refresh()

    def build_aux_batch(
        self,
        batch: InteractionBatch,
        device: torch.device,
    ) -> UnKDAuxBatch:
        if self._epoch_pos_items is None or self._epoch_neg_items is None:
            self.refresh()

        users = batch.unique_users.long().to(device)
        if device.type == "cpu":
            pos_by_user = self._epoch_pos_items
            neg_by_user = self._epoch_neg_items
        else:
            cached = self._device_cache.get(device)
            if cached is None:
                cached = (
                    self._epoch_pos_items.to(device),
                    self._epoch_neg_items.to(device),
                )
                self._device_cache[device] = cached
            pos_by_user, neg_by_user = cached

        return UnKDAuxBatch(
            users=users,
            pos_items=pos_by_user[users],
            neg_items=neg_by_user[users],
        )

    def compute_loss(
        self,
        student: torch.nn.Module,
        batch: InteractionBatch,
        aux_batch: UnKDAuxBatch | None = None,
    ) -> torch.Tensor:
        if aux_batch is None:
            return torch.zeros((), device=batch.users.device)

        if _can_score_unkd_items_together(student):
            items = torch.cat([aux_batch.pos_items, aux_batch.neg_items], dim=1)
            scores = student.score_items(aux_batch.users, items)
            pos_scores, neg_scores = scores.split(
                [aux_batch.pos_items.size(1), aux_batch.neg_items.size(1)],
                dim=1,
            )
        else:
            pos_scores = student.score_items(aux_batch.users, aux_batch.pos_items)
            neg_scores = student.score_items(aux_batch.users, aux_batch.neg_items)
        return self.lambda_unkd * unkd_ranking_loss(pos_scores, neg_scores)

    def refresh(self) -> None:
        if self.dataset is None or not self._ranked_group_items:
            raise RuntimeError("UnKDDistiller must be initialized before sampling.")

        pos_rows: list[torch.Tensor] = []
        neg_rows: list[torch.Tensor] = []
        for user in range(self.dataset.num_users):
            user_pos: list[torch.Tensor] = []
            user_neg: list[torch.Tensor] = []
            for group_idx, sample_count in enumerate(self._group_sample_counts):
                if sample_count <= 0:
                    continue
                ranked_items = self._ranked_group_items[user][group_idx]
                if ranked_items.numel() < 2:
                    continue
                weights = self._ranked_group_weights[user][group_idx]
                pos_items, neg_items = self._sample_ranked_pairs(ranked_items, weights, sample_count)
                user_pos.append(pos_items)
                user_neg.append(neg_items)

            if user_pos:
                pos_row = torch.cat(user_pos, dim=0)
                neg_row = torch.cat(user_neg, dim=0)
            else:
                pos_row, neg_row = self._fallback_user_pairs(user)

            pos_row, neg_row = self._fit_sample_num(pos_row, neg_row)
            pos_rows.append(pos_row)
            neg_rows.append(neg_row)

        self._epoch_pos_items = torch.stack(pos_rows, dim=0).long()
        self._epoch_neg_items = torch.stack(neg_rows, dim=0).long()
        self._device_cache = {}

    def _build_popularity_groups(self, dataset: InteractionDataset) -> list[torch.Tensor]:
        counts = torch.ones(dataset.num_items, dtype=torch.float32)
        for _, item in dataset.interactions:
            counts[int(item)] += 1.0

        sorted_items = torch.argsort(counts, descending=True).tolist()
        target_mass = float(counts.sum().item()) / self.group_count
        groups: list[list[int]] = []
        current: list[int] = []
        current_mass = 0.0
        for item in sorted_items:
            item_mass = float(counts[item].item())
            if current and current_mass + item_mass > target_mass:
                groups.append(current)
                current = []
                current_mass = 0.0
            current.append(int(item))
            current_mass += item_mass
        if current:
            groups.append(current)

        if len(groups) > self.group_count:
            head = groups[: self.group_count - 1]
            tail = [item for group in groups[self.group_count - 1 :] for item in group]
            groups = head + [tail]
        while len(groups) < self.group_count:
            groups.append([])
        return [torch.tensor(group, dtype=torch.long) for group in groups]

    def _build_group_ratios(
        self,
        dataset: InteractionDataset,
        group_items: list[torch.Tensor],
    ) -> list[float]:
        counts = torch.ones(dataset.num_items, dtype=torch.float32)
        for _, item in dataset.interactions:
            counts[int(item)] += 1.0

        avg_popularity = []
        for items in group_items:
            if items.numel() == 0:
                avg_popularity.append(0.0)
            else:
                avg_popularity.append(float(counts[items].mean().item()))

        max_pop = max(avg_popularity) if avg_popularity else 1.0
        min_pop = min(avg_popularity) if avg_popularity else 0.0
        inverse = [max((max_pop + min_pop) - value, 0.0) for value in avg_popularity]
        total = sum(inverse)
        if total <= 0:
            ratios = [1.0 / len(group_items) for _ in group_items]
        else:
            ratios = [max(value / total, 0.1) for value in inverse]

        ratios = [math.pow(value, self.popularity_lambda) for value in ratios]
        total = sum(ratios)
        return [value / total for value in ratios]

    def _allocate_group_samples(self, group_ratios: list[float]) -> list[int]:
        raw = [ratio * self.sample_num for ratio in group_ratios]
        counts = [int(math.floor(value)) for value in raw]
        remaining = self.sample_num - sum(counts)
        order = sorted(range(len(raw)), key=lambda idx: raw[idx] - counts[idx], reverse=True)
        for idx in order[:remaining]:
            counts[idx] += 1
        return counts

    def _build_ranked_group_items(
        self,
        teacher_state: TeacherState,
        dataset: InteractionDataset,
    ) -> list[list[torch.Tensor]]:
        ranked_by_user: list[list[torch.Tensor]] = []
        for user in range(dataset.num_users):
            seen = dataset.seen_items(user)
            user_groups = []
            all_scores = None
            if teacher_state.scorer is not None:
                all_scores = teacher_state.scorer.score_items_for_user(user, teacher_state.num_items).detach().cpu()
            for group_items in self._group_items:
                candidates = [int(item) for item in group_items.tolist() if int(item) not in seen]
                if not candidates:
                    user_groups.append(torch.empty(0, dtype=torch.long))
                    continue
                candidate_tensor = torch.tensor(candidates, dtype=torch.long)
                if all_scores is not None:
                    scores = all_scores[candidate_tensor]
                else:
                    scores = self._teacher_scores(teacher_state, user, candidate_tensor)
                top_k = min(self.rank_top_k, candidate_tensor.numel())
                ranked_idx = torch.topk(scores, k=top_k, dim=0).indices.cpu()
                user_groups.append(candidate_tensor[ranked_idx])
            ranked_by_user.append(user_groups)
        return ranked_by_user

    def _build_ranked_group_weights(
        self,
        ranked_group_items: list[list[torch.Tensor]],
    ) -> list[list[torch.Tensor]]:
        weights_by_user: list[list[torch.Tensor]] = []
        for user_groups in ranked_group_items:
            group_weights = []
            for ranked_items in user_groups:
                rank_count = ranked_items.numel()
                if rank_count == 0:
                    group_weights.append(torch.empty(0, dtype=torch.float32))
                    continue
                ranks = torch.arange(1, rank_count + 1, dtype=torch.float32)
                group_weights.append(torch.exp(-ranks / self.rank_temperature))
            weights_by_user.append(group_weights)
        return weights_by_user

    def _teacher_scores(
        self,
        teacher_state: TeacherState,
        user: int,
        items: torch.Tensor,
    ) -> torch.Tensor:
        if teacher_state.scorer is not None:
            all_scores = teacher_state.scorer.score_items_for_user(user, teacher_state.num_items).detach().cpu()
            return all_scores[items]

        user_emb = teacher_state.user_embeddings[user].detach().cpu()
        item_emb = teacher_state.item_embeddings[items].detach().cpu()
        return torch.matmul(item_emb, user_emb)

    def _sample_ranked_pairs(
        self,
        ranked_items: torch.Tensor,
        weights: torch.Tensor,
        sample_count: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        rank_count = ranked_items.numel()
        first = torch.multinomial(weights, sample_count, replacement=True)
        second = torch.multinomial(weights, sample_count, replacement=True)
        same = first == second
        if same.any() and rank_count > 1:
            second[same] = (second[same] + 1) % rank_count
        better = torch.minimum(first, second)
        worse = torch.maximum(first, second)
        return ranked_items[better], ranked_items[worse]

    def _fallback_user_pairs(self, user: int) -> tuple[torch.Tensor, torch.Tensor]:
        assert self.dataset is not None
        seen = self.dataset.seen_items(user)
        available = [item for item in range(self.dataset.num_items) if item not in seen]
        if len(available) < 2:
            raise ValueError(f"User {user} has fewer than two unseen items for UnKD sampling.")
        items = torch.tensor(available, dtype=torch.long)
        shuffled = items[torch.randperm(items.numel())]
        return shuffled[:1], shuffled[1:2]

    def _fit_sample_num(
        self,
        pos_items: torch.Tensor,
        neg_items: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if pos_items.numel() >= self.sample_num:
            return pos_items[: self.sample_num], neg_items[: self.sample_num]

        repeats = math.ceil(self.sample_num / max(1, pos_items.numel()))
        pos_items = pos_items.repeat(repeats)[: self.sample_num]
        neg_items = neg_items.repeat(repeats)[: self.sample_num]
        return pos_items, neg_items

lambda_unkd = float(lambda_unkd) instance-attribute

sample_num = int(sample_num) instance-attribute

group_count = int(group_count) instance-attribute

popularity_lambda = float(popularity_lambda) instance-attribute

rank_top_k = int(rank_top_k) instance-attribute

rank_temperature = float(rank_temperature) instance-attribute

dataset: InteractionDataset | None = None instance-attribute

teacher_state: TeacherState | None = None instance-attribute

__init__(lambda_unkd: float = 1.0, sample_num: int = 30, group_count: int = 2, popularity_lambda: float = 1.0, rank_top_k: int = 1000, rank_temperature: float = 20.0)

Source code in recdistill/distillers/unkd.py 
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
def __init__(
    self,
    lambda_unkd: float = 1.0,
    sample_num: int = 30,
    group_count: int = 2,
    popularity_lambda: float = 1.0,
    rank_top_k: int = 1000,
    rank_temperature: float = 20.0,
):
    super().__init__()
    if lambda_unkd < 0:
        raise ValueError("lambda_unkd must be >= 0.")
    if sample_num < 1:
        raise ValueError("sample_num must be >= 1.")
    if group_count < 1:
        raise ValueError("group_count must be >= 1.")
    if popularity_lambda < 0:
        raise ValueError("popularity_lambda must be >= 0.")
    if rank_top_k < 1:
        raise ValueError("rank_top_k must be >= 1.")
    if rank_temperature <= 0:
        raise ValueError("rank_temperature must be > 0.")

    self.lambda_unkd = float(lambda_unkd)
    self.sample_num = int(sample_num)
    self.group_count = int(group_count)
    self.popularity_lambda = float(popularity_lambda)
    self.rank_top_k = int(rank_top_k)
    self.rank_temperature = float(rank_temperature)

    self.dataset: InteractionDataset | None = None
    self.teacher_state: TeacherState | None = None
    self._group_items: list[torch.Tensor] = []
    self._group_sample_counts: list[int] = []
    self._ranked_group_items: list[list[torch.Tensor]] = []
    self._ranked_group_weights: list[list[torch.Tensor]] = []
    self._epoch_pos_items: torch.Tensor | None = None
    self._epoch_neg_items: torch.Tensor | None = None
    self._device_cache: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {}

on_train_start(teacher_state: TeacherState, dataset: InteractionDataset) -> None

Source code in recdistill/distillers/unkd.py 
79
80
81
82
83
84
85
86
87
def on_train_start(self, teacher_state: TeacherState, dataset: InteractionDataset) -> None:
    self.teacher_state = teacher_state
    self.dataset = dataset
    self._group_items = self._build_popularity_groups(dataset)
    group_ratios = self._build_group_ratios(dataset, self._group_items)
    self._group_sample_counts = self._allocate_group_samples(group_ratios)
    self._ranked_group_items = self._build_ranked_group_items(teacher_state, dataset)
    self._ranked_group_weights = self._build_ranked_group_weights(self._ranked_group_items)
    self.refresh()

on_epoch_start() -> None

Source code in recdistill/distillers/unkd.py 
89
90
def on_epoch_start(self) -> None:
    self.refresh()

build_aux_batch(batch: InteractionBatch, device: torch.device) -> UnKDAuxBatch

Source code in recdistill/distillers/unkd.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def build_aux_batch(
    self,
    batch: InteractionBatch,
    device: torch.device,
) -> UnKDAuxBatch:
    if self._epoch_pos_items is None or self._epoch_neg_items is None:
        self.refresh()

    users = batch.unique_users.long().to(device)
    if device.type == "cpu":
        pos_by_user = self._epoch_pos_items
        neg_by_user = self._epoch_neg_items
    else:
        cached = self._device_cache.get(device)
        if cached is None:
            cached = (
                self._epoch_pos_items.to(device),
                self._epoch_neg_items.to(device),
            )
            self._device_cache[device] = cached
        pos_by_user, neg_by_user = cached

    return UnKDAuxBatch(
        users=users,
        pos_items=pos_by_user[users],
        neg_items=neg_by_user[users],
    )

compute_loss(student: torch.nn.Module, batch: InteractionBatch, aux_batch: UnKDAuxBatch | None = None) -> torch.Tensor

Source code in recdistill/distillers/unkd.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def compute_loss(
    self,
    student: torch.nn.Module,
    batch: InteractionBatch,
    aux_batch: UnKDAuxBatch | None = None,
) -> torch.Tensor:
    if aux_batch is None:
        return torch.zeros((), device=batch.users.device)

    if _can_score_unkd_items_together(student):
        items = torch.cat([aux_batch.pos_items, aux_batch.neg_items], dim=1)
        scores = student.score_items(aux_batch.users, items)
        pos_scores, neg_scores = scores.split(
            [aux_batch.pos_items.size(1), aux_batch.neg_items.size(1)],
            dim=1,
        )
    else:
        pos_scores = student.score_items(aux_batch.users, aux_batch.pos_items)
        neg_scores = student.score_items(aux_batch.users, aux_batch.neg_items)
    return self.lambda_unkd * unkd_ranking_loss(pos_scores, neg_scores)

refresh() -> None

Source code in recdistill/distillers/unkd.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def refresh(self) -> None:
    if self.dataset is None or not self._ranked_group_items:
        raise RuntimeError("UnKDDistiller must be initialized before sampling.")

    pos_rows: list[torch.Tensor] = []
    neg_rows: list[torch.Tensor] = []
    for user in range(self.dataset.num_users):
        user_pos: list[torch.Tensor] = []
        user_neg: list[torch.Tensor] = []
        for group_idx, sample_count in enumerate(self._group_sample_counts):
            if sample_count <= 0:
                continue
            ranked_items = self._ranked_group_items[user][group_idx]
            if ranked_items.numel() < 2:
                continue
            weights = self._ranked_group_weights[user][group_idx]
            pos_items, neg_items = self._sample_ranked_pairs(ranked_items, weights, sample_count)
            user_pos.append(pos_items)
            user_neg.append(neg_items)

        if user_pos:
            pos_row = torch.cat(user_pos, dim=0)
            neg_row = torch.cat(user_neg, dim=0)
        else:
            pos_row, neg_row = self._fallback_user_pairs(user)

        pos_row, neg_row = self._fit_sample_num(pos_row, neg_row)
        pos_rows.append(pos_row)
        neg_rows.append(neg_row)

    self._epoch_pos_items = torch.stack(pos_rows, dim=0).long()
    self._epoch_neg_items = torch.stack(neg_rows, dim=0).long()
    self._device_cache = {}

unkd_ranking_loss(pos_scores: torch.Tensor, neg_scores: torch.Tensor) -> torch.Tensor

Source code in recdistill/distillers/unkd.py 
14
15
def unkd_ranking_loss(pos_scores: torch.Tensor, neg_scores: torch.Tensor) -> torch.Tensor:
    return -F.logsigmoid(pos_scores - neg_scores).mean()

HTD

GroupMLP

Bases: Module

Source code in recdistill/distillers/htd.py 
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
class GroupMLP(nn.Module):
    def __init__(
        self,
        in_dim: int,
        hidden_dim: int,
        out_dim: int,
        num_groups: int,
    ):
        super().__init__()
        self.num_groups = num_groups
        self.hidden_dim = hidden_dim
        self.fc1 = nn.Linear(in_dim, num_groups * hidden_dim)
        self.fc2 = nn.Conv1d(
            in_channels=num_groups * hidden_dim,
            out_channels=num_groups * out_dim,
            kernel_size=1,
            groups=num_groups,
        )
        self.relu = nn.ReLU()

    def forward(self, x: Tensor) -> Tensor:
        batch_size = x.size(0)
        hidden = self.relu(self.fc1(x))
        hidden = hidden.unsqueeze(-1)
        out = self.fc2(hidden)
        out = out.view(batch_size, self.num_groups, -1)
        return out

num_groups = num_groups instance-attribute

hidden_dim = hidden_dim instance-attribute

fc1 = nn.Linear(in_dim, num_groups * hidden_dim) instance-attribute

fc2 = nn.Conv1d(in_channels=(num_groups * hidden_dim), out_channels=(num_groups * out_dim), kernel_size=1, groups=num_groups) instance-attribute

relu = nn.ReLU() instance-attribute

__init__(in_dim: int, hidden_dim: int, out_dim: int, num_groups: int)

Source code in recdistill/distillers/htd.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def __init__(
    self,
    in_dim: int,
    hidden_dim: int,
    out_dim: int,
    num_groups: int,
):
    super().__init__()
    self.num_groups = num_groups
    self.hidden_dim = hidden_dim
    self.fc1 = nn.Linear(in_dim, num_groups * hidden_dim)
    self.fc2 = nn.Conv1d(
        in_channels=num_groups * hidden_dim,
        out_channels=num_groups * out_dim,
        kernel_size=1,
        groups=num_groups,
    )
    self.relu = nn.ReLU()

forward(x: Tensor) -> Tensor

Source code in recdistill/distillers/htd.py 
33
34
35
36
37
38
39
def forward(self, x: Tensor) -> Tensor:
    batch_size = x.size(0)
    hidden = self.relu(self.fc1(x))
    hidden = hidden.unsqueeze(-1)
    out = self.fc2(hidden)
    out = out.view(batch_size, self.num_groups, -1)
    return out

HTDistiller

Bases: Distiller

Hierarchical Topology Distillation (adapted).

Uses TeacherState for teacher embeddings and the student's get_all_*_embeddings() API to index student embeddings.

Source code in recdistill/distillers/htd.py 
 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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
class HTDistiller(Distiller):
    """Hierarchical Topology Distillation (adapted).

    Uses `TeacherState` for teacher embeddings and the student's
    `get_all_*_embeddings()` API to index student embeddings.
    """

    def __init__(
        self,
        lambda_td: float = 1e-3,
        alpha: float = 0.5,
        num_groups: int = 40,
        topology_mode: str = "group_pe",
        initial_tau: float = 1.0,
        min_tau: float = 1e-10,
        decay_epochs: int = 100,
        entity_sample_size: int = 0,
    ):
        super().__init__()
        self.lambda_td = lambda_td
        self.alpha = alpha
        self.K = num_groups
        if topology_mode not in ["group_pp", "group_pe"]:
            raise ValueError("topology_mode must be 'group_pp' or 'group_pe'")
        self.topology_mode = topology_mode

        self.initial_tau = initial_tau
        self.min_tau = min_tau
        self.decay_epochs = decay_epochs
        self.entity_sample_size = int(entity_sample_size)
        self.tau = initial_tau
        self._epoch = 0

        self.register_buffer("_teacher_users", torch.empty(0), persistent=False)
        self.register_buffer("_teacher_items", torch.empty(0), persistent=False)

        self.v_user = None
        self.v_item = None
        self.f_user = None
        self.f_item = None
        self._teacher_dim = None

    def on_train_start(self, teacher_state: TeacherState, dataset) -> None:
        del dataset
        device = teacher_state.device
        self._teacher_users = teacher_state.user_embeddings.detach().to(device)
        self._teacher_items = teacher_state.item_embeddings.detach().to(device)

        d_t = teacher_state.embedding_dim
        self._teacher_dim = d_t

        self.v_user = nn.Sequential(nn.Linear(d_t, self.K), nn.Softmax(dim=1))
        self.v_item = nn.Sequential(nn.Linear(d_t, self.K), nn.Softmax(dim=1))

        self.to(device)

    def _ensure_student_adapters(self, student_dim: int, device: torch.device) -> None:
        if self._teacher_dim is None:
            raise RuntimeError("HTDistiller is not initialized. Call on_train_start first.")

        needs_init = self.f_user is None or self.f_item is None
        if not needs_init:
            user_in = int(self.f_user.fc1.in_features)
            item_in = int(self.f_item.fc1.in_features)
            needs_init = user_in != student_dim or item_in != student_dim

        if needs_init:
            hidden_dim = max(1, (student_dim + self._teacher_dim) // 2)
            self.f_user = GroupMLP(student_dim, hidden_dim, self._teacher_dim, self.K).to(device)
            self.f_item = GroupMLP(student_dim, hidden_dim, self._teacher_dim, self.K).to(device)

    def on_epoch_start(self) -> None:
        self._update_temperature(self._epoch)
        self._epoch += 1

    def _update_temperature(self, current_epoch: int) -> None:
        if current_epoch >= self.decay_epochs:
            self.tau = self.min_tau
        else:
            ratio = current_epoch / self.decay_epochs
            self.tau = self.initial_tau * ((self.min_tau / self.initial_tau) ** ratio)
        self.tau = max(self.tau, self.min_tau)

    def _cosine_similarity(self, x: Tensor, y: Tensor) -> Tensor:
        return F.normalize(x, dim=-1, eps=1e-8) @ F.normalize(y, dim=-1, eps=1e-8).T

    def _sample_entities(self, users: torch.Tensor, items: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if self.entity_sample_size <= 0:
            return users, items

        total = int(users.numel() + items.numel())
        if total <= self.entity_sample_size:
            return users, items

        user_count = int(users.numel())
        item_count = int(items.numel())
        keep_users = min(user_count, max(1, round(self.entity_sample_size * user_count / total)))
        keep_items = min(item_count, max(1, self.entity_sample_size - keep_users))

        if keep_users + keep_items > self.entity_sample_size:
            keep_items = max(1, self.entity_sample_size - keep_users)

        users = users[torch.randperm(user_count, device=users.device)[:keep_users]]
        items = items[torch.randperm(item_count, device=items.device)[:keep_items]]
        return users, items

    def _compute_ga_loss(self, t_emb: Tensor, s_emb: Tensor, v_net: nn.Module, f_net: nn.Module):
        alpha = v_net(t_emb) + 1e-10
        z_soft = F.gumbel_softmax(alpha.log(), tau=self.tau, hard=False, dim=-1)
        expert_outputs = f_net(s_emb)
        recon = torch.sum(expert_outputs * z_soft.unsqueeze(-1), dim=1)
        loss = torch.sum((t_emb - recon) ** 2, dim=-1).sum()
        return loss, z_soft

    def compute_loss(self, student: nn.Module, batch: InteractionBatch, aux_batch: object | None = None) -> Tensor:
        del aux_batch
        u_indices = batch.unique_users
        i_indices = batch.unique_items
        u_indices, i_indices = self._sample_entities(u_indices, i_indices)

        device = self._teacher_users.device if self._teacher_users.numel() else next(self.parameters()).device

        t_u = self._teacher_users[u_indices].to(device)
        t_i = self._teacher_items[i_indices].to(device)

        s_all_u = student.get_all_user_embeddings().to(device)
        s_all_i = student.get_all_item_embeddings().to(device)
        self._ensure_student_adapters(student_dim=int(s_all_u.size(1)), device=device)
        s_u = s_all_u[u_indices]
        s_i = s_all_i[i_indices]

        loss_user, z_soft_u = self._compute_ga_loss(t_u, s_u, self.v_user, self.f_user)
        loss_item, z_soft_i = self._compute_ga_loss(t_i, s_i, self.v_item, self.f_item)
        L_GA = loss_user + loss_item

        E_t = torch.cat([t_u, t_i], dim=0)
        E_s = torch.cat([s_u, s_i], dim=0)

        with torch.no_grad():
            z_u = self.v_user(t_u).argmax(dim=1)
            z_i = self.v_item(t_i).argmax(dim=1) + self.K
            z_all = torch.cat([z_u, z_i], dim=0)

            ones = torch.ones_like(z_all, dtype=torch.float).unsqueeze(1)
            counts = torch.zeros(2 * self.K, 1, device=z_all.device)
            counts.scatter_add_(0, z_all.unsqueeze(1), ones)
            counts = counts + 1e-10

            P_t_sum = torch.zeros(2 * self.K, E_t.shape[1], device=E_t.device)
            P_s_sum = torch.zeros(2 * self.K, E_s.shape[1], device=E_s.device)
            P_t_sum.index_add_(0, z_all, E_t)
            P_s_sum.index_add_(0, z_all, E_s)

            active_groups = z_all.unique()
            P_t = (P_t_sum / counts)[active_groups]
            P_s = (P_s_sum / counts)[active_groups]

            Z = F.one_hot(z_all, num_classes=2 * self.K).float()
            M = Z @ Z.T

        sim_tt = self._cosine_similarity(E_t, E_t) * M
        sim_ss = self._cosine_similarity(E_s, E_s) * M

        valid_mask = sim_tt > 0.0
        sim_tt_filtered = sim_tt[valid_mask]
        sim_ss_filtered = sim_ss[valid_mask]
        L_entity = torch.sum((sim_tt_filtered - sim_ss_filtered) ** 2)

        if self.topology_mode == "group_pp":
            proto_tt = self._cosine_similarity(P_t, P_t).view(-1)
            proto_ss = self._cosine_similarity(P_s, P_s).view(-1)
            L_group = torch.sum((proto_tt - proto_ss) ** 2)
        else:
            proto_dist_t = self._cosine_similarity(P_t, E_t).view(-1)
            proto_dist_s = self._cosine_similarity(P_s, E_s).view(-1)
            L_group = torch.sum((proto_dist_t - proto_dist_s) ** 2)

        L_TD = L_entity + L_group
        HTD_loss = L_TD * self.alpha + L_GA * (1 - self.alpha)
        normalizer = max(1, int(batch.users.numel()))
        return (HTD_loss / normalizer) * self.lambda_td

lambda_td = lambda_td instance-attribute

alpha = alpha instance-attribute

K = num_groups instance-attribute

topology_mode = topology_mode instance-attribute

initial_tau = initial_tau instance-attribute

min_tau = min_tau instance-attribute

decay_epochs = decay_epochs instance-attribute

entity_sample_size = int(entity_sample_size) instance-attribute

tau = initial_tau instance-attribute

v_user = None instance-attribute

v_item = None instance-attribute

f_user = None instance-attribute

f_item = None instance-attribute

__init__(lambda_td: float = 0.001, alpha: float = 0.5, num_groups: int = 40, topology_mode: str = 'group_pe', initial_tau: float = 1.0, min_tau: float = 1e-10, decay_epochs: int = 100, entity_sample_size: int = 0)

Source code in recdistill/distillers/htd.py 
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
def __init__(
    self,
    lambda_td: float = 1e-3,
    alpha: float = 0.5,
    num_groups: int = 40,
    topology_mode: str = "group_pe",
    initial_tau: float = 1.0,
    min_tau: float = 1e-10,
    decay_epochs: int = 100,
    entity_sample_size: int = 0,
):
    super().__init__()
    self.lambda_td = lambda_td
    self.alpha = alpha
    self.K = num_groups
    if topology_mode not in ["group_pp", "group_pe"]:
        raise ValueError("topology_mode must be 'group_pp' or 'group_pe'")
    self.topology_mode = topology_mode

    self.initial_tau = initial_tau
    self.min_tau = min_tau
    self.decay_epochs = decay_epochs
    self.entity_sample_size = int(entity_sample_size)
    self.tau = initial_tau
    self._epoch = 0

    self.register_buffer("_teacher_users", torch.empty(0), persistent=False)
    self.register_buffer("_teacher_items", torch.empty(0), persistent=False)

    self.v_user = None
    self.v_item = None
    self.f_user = None
    self.f_item = None
    self._teacher_dim = None

on_train_start(teacher_state: TeacherState, dataset) -> None

Source code in recdistill/distillers/htd.py 
84
85
86
87
88
89
90
91
92
93
94
95
96
def on_train_start(self, teacher_state: TeacherState, dataset) -> None:
    del dataset
    device = teacher_state.device
    self._teacher_users = teacher_state.user_embeddings.detach().to(device)
    self._teacher_items = teacher_state.item_embeddings.detach().to(device)

    d_t = teacher_state.embedding_dim
    self._teacher_dim = d_t

    self.v_user = nn.Sequential(nn.Linear(d_t, self.K), nn.Softmax(dim=1))
    self.v_item = nn.Sequential(nn.Linear(d_t, self.K), nn.Softmax(dim=1))

    self.to(device)

on_epoch_start() -> None

Source code in recdistill/distillers/htd.py 
113
114
115
def on_epoch_start(self) -> None:
    self._update_temperature(self._epoch)
    self._epoch += 1

compute_loss(student: nn.Module, batch: InteractionBatch, aux_batch: object | None = None) -> Tensor

Source code in recdistill/distillers/htd.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
def compute_loss(self, student: nn.Module, batch: InteractionBatch, aux_batch: object | None = None) -> Tensor:
    del aux_batch
    u_indices = batch.unique_users
    i_indices = batch.unique_items
    u_indices, i_indices = self._sample_entities(u_indices, i_indices)

    device = self._teacher_users.device if self._teacher_users.numel() else next(self.parameters()).device

    t_u = self._teacher_users[u_indices].to(device)
    t_i = self._teacher_items[i_indices].to(device)

    s_all_u = student.get_all_user_embeddings().to(device)
    s_all_i = student.get_all_item_embeddings().to(device)
    self._ensure_student_adapters(student_dim=int(s_all_u.size(1)), device=device)
    s_u = s_all_u[u_indices]
    s_i = s_all_i[i_indices]

    loss_user, z_soft_u = self._compute_ga_loss(t_u, s_u, self.v_user, self.f_user)
    loss_item, z_soft_i = self._compute_ga_loss(t_i, s_i, self.v_item, self.f_item)
    L_GA = loss_user + loss_item

    E_t = torch.cat([t_u, t_i], dim=0)
    E_s = torch.cat([s_u, s_i], dim=0)

    with torch.no_grad():
        z_u = self.v_user(t_u).argmax(dim=1)
        z_i = self.v_item(t_i).argmax(dim=1) + self.K
        z_all = torch.cat([z_u, z_i], dim=0)

        ones = torch.ones_like(z_all, dtype=torch.float).unsqueeze(1)
        counts = torch.zeros(2 * self.K, 1, device=z_all.device)
        counts.scatter_add_(0, z_all.unsqueeze(1), ones)
        counts = counts + 1e-10

        P_t_sum = torch.zeros(2 * self.K, E_t.shape[1], device=E_t.device)
        P_s_sum = torch.zeros(2 * self.K, E_s.shape[1], device=E_s.device)
        P_t_sum.index_add_(0, z_all, E_t)
        P_s_sum.index_add_(0, z_all, E_s)

        active_groups = z_all.unique()
        P_t = (P_t_sum / counts)[active_groups]
        P_s = (P_s_sum / counts)[active_groups]

        Z = F.one_hot(z_all, num_classes=2 * self.K).float()
        M = Z @ Z.T

    sim_tt = self._cosine_similarity(E_t, E_t) * M
    sim_ss = self._cosine_similarity(E_s, E_s) * M

    valid_mask = sim_tt > 0.0
    sim_tt_filtered = sim_tt[valid_mask]
    sim_ss_filtered = sim_ss[valid_mask]
    L_entity = torch.sum((sim_tt_filtered - sim_ss_filtered) ** 2)

    if self.topology_mode == "group_pp":
        proto_tt = self._cosine_similarity(P_t, P_t).view(-1)
        proto_ss = self._cosine_similarity(P_s, P_s).view(-1)
        L_group = torch.sum((proto_tt - proto_ss) ** 2)
    else:
        proto_dist_t = self._cosine_similarity(P_t, E_t).view(-1)
        proto_dist_s = self._cosine_similarity(P_s, E_s).view(-1)
        L_group = torch.sum((proto_dist_t - proto_dist_s) ** 2)

    L_TD = L_entity + L_group
    HTD_loss = L_TD * self.alpha + L_GA * (1 - self.alpha)
    normalizer = max(1, int(batch.users.numel()))
    return (HTD_loss / normalizer) * self.lambda_td

FTD

FTDistiller

Bases: Distiller

Source code in recdistill/distillers/ftd.py 
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
class FTDistiller(Distiller):
    def __init__(self, lambda_td: float = 1e-3, entity_sample_size: int = 0):
        super().__init__()
        self.lambda_td = lambda_td
        self.entity_sample_size = int(entity_sample_size)
        self.register_buffer("_teacher_users", torch.empty(0), persistent=False)
        self.register_buffer("_teacher_items", torch.empty(0), persistent=False)

    def on_train_start(self, teacher_state: TeacherState, dataset) -> None:
        del dataset
        device = teacher_state.device
        self._teacher_users = teacher_state.user_embeddings.detach().to(device)
        self._teacher_items = teacher_state.item_embeddings.detach().to(device)

    def _cosine_similarity(self, x: Tensor, y: torch.Tensor) -> Tensor:
        x_norm = F.normalize(x, dim=-1, eps=1e-8)
        y_norm = F.normalize(y, dim=-1, eps=1e-8)
        return x_norm @ y_norm.T

    def _sample_entities(self, users: torch.Tensor, items: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if self.entity_sample_size <= 0:
            return users, items

        total = int(users.numel() + items.numel())
        if total <= self.entity_sample_size:
            return users, items

        user_count = int(users.numel())
        item_count = int(items.numel())
        keep_users = min(user_count, max(1, round(self.entity_sample_size * user_count / total)))
        keep_items = min(item_count, max(1, self.entity_sample_size - keep_users))

        if keep_users + keep_items > self.entity_sample_size:
            keep_items = max(1, self.entity_sample_size - keep_users)

        users = users[torch.randperm(user_count, device=users.device)[:keep_users]]
        items = items[torch.randperm(item_count, device=items.device)[:keep_items]]
        return users, items

    def compute_loss(self, student: nn.Module, batch: InteractionBatch, aux_batch: object | None = None) -> Tensor:
        del aux_batch
        u_indices = batch.unique_users
        i_indices = batch.unique_items
        u_indices, i_indices = self._sample_entities(u_indices, i_indices)

        device = self._teacher_users.device if self._teacher_users.numel() else next(self.parameters()).device
        t_u = self._teacher_users[u_indices].to(device)
        t_i = self._teacher_items[i_indices].to(device)

        s_all_u = student.get_all_user_embeddings().to(device)
        s_all_i = student.get_all_item_embeddings().to(device)
        s_u = s_all_u[u_indices]
        s_i = s_all_i[i_indices]

        E_t = torch.cat([t_u, t_i], dim=0)
        E_s = torch.cat([s_u, s_i], dim=0)

        A_t = self._cosine_similarity(E_t, E_t)
        A_s = self._cosine_similarity(E_s, E_s)
        L_FTD = torch.sum((A_t - A_s) ** 2)
        normalizer = max(1, int(batch.users.numel()))
        return (L_FTD / normalizer) * self.lambda_td

lambda_td = lambda_td instance-attribute

entity_sample_size = int(entity_sample_size) instance-attribute

__init__(lambda_td: float = 0.001, entity_sample_size: int = 0)

Source code in recdistill/distillers/ftd.py 
14
15
16
17
18
19
def __init__(self, lambda_td: float = 1e-3, entity_sample_size: int = 0):
    super().__init__()
    self.lambda_td = lambda_td
    self.entity_sample_size = int(entity_sample_size)
    self.register_buffer("_teacher_users", torch.empty(0), persistent=False)
    self.register_buffer("_teacher_items", torch.empty(0), persistent=False)

on_train_start(teacher_state: TeacherState, dataset) -> None

Source code in recdistill/distillers/ftd.py 
21
22
23
24
25
def on_train_start(self, teacher_state: TeacherState, dataset) -> None:
    del dataset
    device = teacher_state.device
    self._teacher_users = teacher_state.user_embeddings.detach().to(device)
    self._teacher_items = teacher_state.item_embeddings.detach().to(device)

compute_loss(student: nn.Module, batch: InteractionBatch, aux_batch: object | None = None) -> Tensor

Source code in recdistill/distillers/ftd.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def compute_loss(self, student: nn.Module, batch: InteractionBatch, aux_batch: object | None = None) -> Tensor:
    del aux_batch
    u_indices = batch.unique_users
    i_indices = batch.unique_items
    u_indices, i_indices = self._sample_entities(u_indices, i_indices)

    device = self._teacher_users.device if self._teacher_users.numel() else next(self.parameters()).device
    t_u = self._teacher_users[u_indices].to(device)
    t_i = self._teacher_items[i_indices].to(device)

    s_all_u = student.get_all_user_embeddings().to(device)
    s_all_i = student.get_all_item_embeddings().to(device)
    s_u = s_all_u[u_indices]
    s_i = s_all_i[i_indices]

    E_t = torch.cat([t_u, t_i], dim=0)
    E_s = torch.cat([s_u, s_i], dim=0)

    A_t = self._cosine_similarity(E_t, E_t)
    A_s = self._cosine_similarity(E_s, E_s)
    L_FTD = torch.sum((A_t - A_s) ** 2)
    normalizer = max(1, int(batch.users.numel()))
    return (L_FTD / normalizer) * self.lambda_td

Composite Distillation

CompositeDistiller

Bases: Distiller

Source code in recdistill/distillers/composite.py 
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
class CompositeDistiller(Distiller):
    def __init__(self, distillers: list[Distiller]):
        super().__init__()
        self.distillers = torch.nn.ModuleList(distillers)

    @staticmethod
    def _distiller_key(distiller: Distiller, index: int) -> str:
        return f"{distiller.__class__.__name__}:{index}"

    def on_train_start(self, teacher_state: TeacherState, dataset: InteractionDataset) -> None:
        for distiller in self.distillers:
            distiller.on_train_start(teacher_state, dataset)

    def on_epoch_start(self) -> None:
        for distiller in self.distillers:
            distiller.on_epoch_start()

    def build_aux_batch(self, batch: InteractionBatch, device: torch.device) -> dict[str, object]:
        aux_batches: dict[str, object] = {}
        for index, distiller in enumerate(self.distillers):
            aux = distiller.build_aux_batch(batch, device)
            if aux is not None:
                aux_batches[self._distiller_key(distiller, index)] = aux
        return aux_batches

    def compute_loss(
        self,
        student: torch.nn.Module,
        batch: InteractionBatch,
        aux_batch: dict[str, object] | None = None,
    ) -> torch.Tensor:
        if not self.distillers:
            return torch.zeros((), device=batch.users.device)

        total_loss = torch.zeros((), device=batch.users.device)
        aux_batch = aux_batch or {}
        for index, distiller in enumerate(self.distillers):
            total_loss = total_loss + distiller.compute_loss(
                student,
                batch,
                aux_batch.get(self._distiller_key(distiller, index)),
            )
        return total_loss

distillers = torch.nn.ModuleList(distillers) instance-attribute

__init__(distillers: list[Distiller])

Source code in recdistill/distillers/composite.py 
12
13
14
def __init__(self, distillers: list[Distiller]):
    super().__init__()
    self.distillers = torch.nn.ModuleList(distillers)

on_train_start(teacher_state: TeacherState, dataset: InteractionDataset) -> None

Source code in recdistill/distillers/composite.py 
20
21
22
def on_train_start(self, teacher_state: TeacherState, dataset: InteractionDataset) -> None:
    for distiller in self.distillers:
        distiller.on_train_start(teacher_state, dataset)

on_epoch_start() -> None

Source code in recdistill/distillers/composite.py 
24
25
26
def on_epoch_start(self) -> None:
    for distiller in self.distillers:
        distiller.on_epoch_start()

build_aux_batch(batch: InteractionBatch, device: torch.device) -> dict[str, object]

Source code in recdistill/distillers/composite.py 
28
29
30
31
32
33
34
def build_aux_batch(self, batch: InteractionBatch, device: torch.device) -> dict[str, object]:
    aux_batches: dict[str, object] = {}
    for index, distiller in enumerate(self.distillers):
        aux = distiller.build_aux_batch(batch, device)
        if aux is not None:
            aux_batches[self._distiller_key(distiller, index)] = aux
    return aux_batches

compute_loss(student: torch.nn.Module, batch: InteractionBatch, aux_batch: dict[str, object] | None = None) -> torch.Tensor

Source code in recdistill/distillers/composite.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def compute_loss(
    self,
    student: torch.nn.Module,
    batch: InteractionBatch,
    aux_batch: dict[str, object] | None = None,
) -> torch.Tensor:
    if not self.distillers:
        return torch.zeros((), device=batch.users.device)

    total_loss = torch.zeros((), device=batch.users.device)
    aux_batch = aux_batch or {}
    for index, distiller in enumerate(self.distillers):
        total_loss = total_loss + distiller.compute_loss(
            student,
            batch,
            aux_batch.get(self._distiller_key(distiller, index)),
        )
    return total_loss

Distillation Samplers

AuxiliarySampler

Bases: ABC

Source code in recdistill/samplers/base.py 
 8
 9
10
11
12
13
14
15
16
17
class AuxiliarySampler(ABC):
    def initialize(self, dataset, teacher_state) -> None:
        del dataset, teacher_state

    def refresh(self) -> None:
        return None

    @abstractmethod
    def sample(self, indices: torch.Tensor, device: torch.device) -> object:
        raise NotImplementedError

initialize(dataset, teacher_state) -> None

Source code in recdistill/samplers/base.py 
 9
10
def initialize(self, dataset, teacher_state) -> None:
    del dataset, teacher_state

refresh() -> None

Source code in recdistill/samplers/base.py 
12
13
def refresh(self) -> None:
    return None

sample(indices: torch.Tensor, device: torch.device) -> object abstractmethod

Source code in recdistill/samplers/base.py 
15
16
17
@abstractmethod
def sample(self, indices: torch.Tensor, device: torch.device) -> object:
    raise NotImplementedError

BPRNegativeSampler

Source code in recdistill/samplers/negative.py 
 8
 9
10
11
12
13
14
15
16
17
class BPRNegativeSampler:
    def __init__(self, dataset: InteractionDataset):
        self.dataset = dataset

    def sample(self, user: int) -> int:
        seen = self.dataset.seen_items(user)
        while True:
            item = random.randrange(self.dataset.num_items)
            if item not in seen:
                return item

dataset = dataset instance-attribute

__init__(dataset: InteractionDataset)

Source code in recdistill/samplers/negative.py 
 9
10
def __init__(self, dataset: InteractionDataset):
    self.dataset = dataset

sample(user: int) -> int

Source code in recdistill/samplers/negative.py 
12
13
14
15
16
17
def sample(self, user: int) -> int:
    seen = self.dataset.seen_items(user)
    while True:
        item = random.randrange(self.dataset.num_items)
        if item not in seen:
            return item

RRDSampler

Bases: AuxiliarySampler

Source code in recdistill/samplers/rrd.py 
 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
116
117
118
class RRDSampler(AuxiliarySampler):
    def __init__(
        self,
        interesting_size: int,
        uninteresting_size: int,
        temperature: float,
        teacher_topk: dict[int, list[int]] | None = None,
        topk_provider: TeacherTopKProvider | None = None,
    ):
        self.interesting_size = interesting_size
        self.uninteresting_size = uninteresting_size
        self.temperature = temperature
        self.teacher_topk = teacher_topk
        self.topk_provider = topk_provider or TeacherTopKProvider(top_k=500)
        self.dataset: InteractionDataset | None = None
        self._teacher_topk_tensors: dict[int, torch.Tensor] = {}
        self._ranking_weights: dict[int, torch.Tensor] = {}
        self._available_uninteresting: dict[int, torch.Tensor] = {}
        self._epoch_interesting: torch.Tensor | None = None
        self._epoch_uninteresting: torch.Tensor | None = None
        self._device_cache: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {}

    def initialize(self, dataset: InteractionDataset, teacher_state: TeacherState) -> None:
        self.dataset = dataset
        if self.teacher_topk is None:
            self.teacher_topk = self.topk_provider.build(teacher_state=teacher_state, dataset=dataset)

        all_items = set(range(dataset.num_items))
        self._teacher_topk_tensors = {}
        self._ranking_weights = {}
        self._available_uninteresting = {}
        self._epoch_interesting = None
        self._epoch_uninteresting = None
        self._device_cache = {}
        for user in range(dataset.num_users):
            topk_items = self.teacher_topk.get(user, [])
            if len(topk_items) < self.interesting_size:
                raise ValueError(
                    f"User {user} has only {len(topk_items)} teacher top-k items, "
                    f"but RRDSampler requires at least {self.interesting_size}."
                )
            self._teacher_topk_tensors[user] = torch.tensor(topk_items, dtype=torch.long)
            weights = [math.exp(-(idx + 1) / self.temperature) for idx in range(len(topk_items))]
            self._ranking_weights[user] = torch.tensor(weights, dtype=torch.float32)

            blocked = set(topk_items) | dataset.seen_items(user)
            available = sorted(all_items - blocked)
            if len(available) < self.uninteresting_size:
                raise ValueError(
                    f"User {user} has only {len(available)} uninteresting items available, "
                    f"but RRDSampler requires at least {self.uninteresting_size}."
                )
            self._available_uninteresting[user] = torch.tensor(available, dtype=torch.long)

    def refresh(self) -> None:
        if self.dataset is None or self.teacher_topk is None or not self._ranking_weights:
            raise RuntimeError("RRDSampler must be initialized before sampling.")

        interesting_rows = []
        uninteresting_rows = []

        for user in range(self.dataset.num_users):
            sampled_positions = torch.multinomial(
                self._ranking_weights[user],
                self.interesting_size,
                replacement=False,
            )
            sampled_positions = sampled_positions.sort().values
            interesting_rows.append(self._teacher_topk_tensors[user][sampled_positions])

            available = self._available_uninteresting[user]
            sampled_uninteresting_idx = torch.randperm(len(available))[: self.uninteresting_size]
            uninteresting_rows.append(available[sampled_uninteresting_idx])

        self._epoch_interesting = torch.stack(interesting_rows, dim=0)
        self._epoch_uninteresting = torch.stack(uninteresting_rows, dim=0)
        self._device_cache = {}

    def sample(self, indices: torch.Tensor, device: torch.device) -> RRDAuxBatch:
        if self.dataset is None or self.teacher_topk is None or not self._ranking_weights:
            raise RuntimeError("RRDSampler must be initialized before sampling.")
        if self._epoch_interesting is None or self._epoch_uninteresting is None:
            self.refresh()

        users = indices.long().to(device)
        if device.type == "cpu":
            interesting_by_user = self._epoch_interesting
            uninteresting_by_user = self._epoch_uninteresting
        else:
            cached = self._device_cache.get(device)
            if cached is None:
                cached = (
                    self._epoch_interesting.to(device),
                    self._epoch_uninteresting.to(device),
                )
                self._device_cache[device] = cached
            interesting_by_user, uninteresting_by_user = cached

        interesting_items = interesting_by_user[users]
        uninteresting_items = uninteresting_by_user[users]
        return RRDAuxBatch(
            users=users,
            interesting_items=interesting_items,
            uninteresting_items=uninteresting_items,
        )

interesting_size = interesting_size instance-attribute

uninteresting_size = uninteresting_size instance-attribute

temperature = temperature instance-attribute

teacher_topk = teacher_topk instance-attribute

topk_provider = topk_provider or TeacherTopKProvider(top_k=500) instance-attribute

dataset: InteractionDataset | None = None instance-attribute

__init__(interesting_size: int, uninteresting_size: int, temperature: float, teacher_topk: dict[int, list[int]] | None = None, topk_provider: TeacherTopKProvider | None = None)

Source code in recdistill/samplers/rrd.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def __init__(
    self,
    interesting_size: int,
    uninteresting_size: int,
    temperature: float,
    teacher_topk: dict[int, list[int]] | None = None,
    topk_provider: TeacherTopKProvider | None = None,
):
    self.interesting_size = interesting_size
    self.uninteresting_size = uninteresting_size
    self.temperature = temperature
    self.teacher_topk = teacher_topk
    self.topk_provider = topk_provider or TeacherTopKProvider(top_k=500)
    self.dataset: InteractionDataset | None = None
    self._teacher_topk_tensors: dict[int, torch.Tensor] = {}
    self._ranking_weights: dict[int, torch.Tensor] = {}
    self._available_uninteresting: dict[int, torch.Tensor] = {}
    self._epoch_interesting: torch.Tensor | None = None
    self._epoch_uninteresting: torch.Tensor | None = None
    self._device_cache: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {}

initialize(dataset: InteractionDataset, teacher_state: TeacherState) -> None

Source code in recdistill/samplers/rrd.py 
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
def initialize(self, dataset: InteractionDataset, teacher_state: TeacherState) -> None:
    self.dataset = dataset
    if self.teacher_topk is None:
        self.teacher_topk = self.topk_provider.build(teacher_state=teacher_state, dataset=dataset)

    all_items = set(range(dataset.num_items))
    self._teacher_topk_tensors = {}
    self._ranking_weights = {}
    self._available_uninteresting = {}
    self._epoch_interesting = None
    self._epoch_uninteresting = None
    self._device_cache = {}
    for user in range(dataset.num_users):
        topk_items = self.teacher_topk.get(user, [])
        if len(topk_items) < self.interesting_size:
            raise ValueError(
                f"User {user} has only {len(topk_items)} teacher top-k items, "
                f"but RRDSampler requires at least {self.interesting_size}."
            )
        self._teacher_topk_tensors[user] = torch.tensor(topk_items, dtype=torch.long)
        weights = [math.exp(-(idx + 1) / self.temperature) for idx in range(len(topk_items))]
        self._ranking_weights[user] = torch.tensor(weights, dtype=torch.float32)

        blocked = set(topk_items) | dataset.seen_items(user)
        available = sorted(all_items - blocked)
        if len(available) < self.uninteresting_size:
            raise ValueError(
                f"User {user} has only {len(available)} uninteresting items available, "
                f"but RRDSampler requires at least {self.uninteresting_size}."
            )
        self._available_uninteresting[user] = torch.tensor(available, dtype=torch.long)

refresh() -> None

Source code in recdistill/samplers/rrd.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def refresh(self) -> None:
    if self.dataset is None or self.teacher_topk is None or not self._ranking_weights:
        raise RuntimeError("RRDSampler must be initialized before sampling.")

    interesting_rows = []
    uninteresting_rows = []

    for user in range(self.dataset.num_users):
        sampled_positions = torch.multinomial(
            self._ranking_weights[user],
            self.interesting_size,
            replacement=False,
        )
        sampled_positions = sampled_positions.sort().values
        interesting_rows.append(self._teacher_topk_tensors[user][sampled_positions])

        available = self._available_uninteresting[user]
        sampled_uninteresting_idx = torch.randperm(len(available))[: self.uninteresting_size]
        uninteresting_rows.append(available[sampled_uninteresting_idx])

    self._epoch_interesting = torch.stack(interesting_rows, dim=0)
    self._epoch_uninteresting = torch.stack(uninteresting_rows, dim=0)
    self._device_cache = {}

sample(indices: torch.Tensor, device: torch.device) -> RRDAuxBatch

Source code in recdistill/samplers/rrd.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def sample(self, indices: torch.Tensor, device: torch.device) -> RRDAuxBatch:
    if self.dataset is None or self.teacher_topk is None or not self._ranking_weights:
        raise RuntimeError("RRDSampler must be initialized before sampling.")
    if self._epoch_interesting is None or self._epoch_uninteresting is None:
        self.refresh()

    users = indices.long().to(device)
    if device.type == "cpu":
        interesting_by_user = self._epoch_interesting
        uninteresting_by_user = self._epoch_uninteresting
    else:
        cached = self._device_cache.get(device)
        if cached is None:
            cached = (
                self._epoch_interesting.to(device),
                self._epoch_uninteresting.to(device),
            )
            self._device_cache[device] = cached
        interesting_by_user, uninteresting_by_user = cached

    interesting_items = interesting_by_user[users]
    uninteresting_items = uninteresting_by_user[users]
    return RRDAuxBatch(
        users=users,
        interesting_items=interesting_items,
        uninteresting_items=uninteresting_items,
    )

TeacherTopKProvider

Source code in recdistill/samplers/teacher_topk.py 
 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
class TeacherTopKProvider:
    def __init__(self, top_k: int):
        self.top_k = top_k

    def build(
        self,
        teacher_state: TeacherState,
        dataset: InteractionDataset,
    ) -> dict[int, list[int]]:
        if teacher_state.scorer is not None:
            return self._build_with_exact_scorer(teacher_state=teacher_state, dataset=dataset)
        if not teacher_state.has_embeddings:
            raise ValueError("TeacherTopKProvider requires teacher embeddings or a scorer.")

        user_emb = teacher_state.user_embeddings
        item_emb = teacher_state.item_embeddings
        scores = torch.matmul(user_emb, item_emb.T)
        topk_by_user: dict[int, list[int]] = {}
        num_teacher_users = int(user_emb.size(0))
        num_teacher_items = int(item_emb.size(0))

        for user in range(min(dataset.num_users, num_teacher_users)):
            seen = sorted(item for item in dataset.seen_items(user) if 0 <= item < num_teacher_items)
            user_scores = scores[user].clone()
            if seen:
                user_scores[seen] = -1e9
            k = min(self.top_k, num_teacher_items - len(seen))
            if k <= 0:
                topk_by_user[user] = []
                continue
            top_items = torch.topk(user_scores, k=k, dim=0).indices.tolist()
            topk_by_user[user] = [int(item) for item in top_items]

        return topk_by_user

    def _build_with_exact_scorer(
        self,
        teacher_state: TeacherState,
        dataset: InteractionDataset,
    ) -> dict[int, list[int]]:
        topk_by_user: dict[int, list[int]] = {}
        num_teacher_users = teacher_state.num_users
        num_teacher_items = teacher_state.num_items

        for user in range(min(dataset.num_users, num_teacher_users)):
            seen = sorted(item for item in dataset.seen_items(user) if 0 <= item < num_teacher_items)
            user_scores = teacher_state.scorer.score_items_for_user(user, num_teacher_items).clone()
            if seen:
                user_scores[seen] = -1e9
            k = min(self.top_k, num_teacher_items - len(seen))
            if k <= 0:
                topk_by_user[user] = []
                continue
            top_items = torch.topk(user_scores, k=k, dim=0).indices.tolist()
            topk_by_user[user] = [int(item) for item in top_items]

        return topk_by_user

top_k = top_k instance-attribute

__init__(top_k: int)

Source code in recdistill/samplers/teacher_topk.py 
10
11
def __init__(self, top_k: int):
    self.top_k = top_k

build(teacher_state: TeacherState, dataset: InteractionDataset) -> dict[int, list[int]]

Source code in recdistill/samplers/teacher_topk.py 
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
def build(
    self,
    teacher_state: TeacherState,
    dataset: InteractionDataset,
) -> dict[int, list[int]]:
    if teacher_state.scorer is not None:
        return self._build_with_exact_scorer(teacher_state=teacher_state, dataset=dataset)
    if not teacher_state.has_embeddings:
        raise ValueError("TeacherTopKProvider requires teacher embeddings or a scorer.")

    user_emb = teacher_state.user_embeddings
    item_emb = teacher_state.item_embeddings
    scores = torch.matmul(user_emb, item_emb.T)
    topk_by_user: dict[int, list[int]] = {}
    num_teacher_users = int(user_emb.size(0))
    num_teacher_items = int(item_emb.size(0))

    for user in range(min(dataset.num_users, num_teacher_users)):
        seen = sorted(item for item in dataset.seen_items(user) if 0 <= item < num_teacher_items)
        user_scores = scores[user].clone()
        if seen:
            user_scores[seen] = -1e9
        k = min(self.top_k, num_teacher_items - len(seen))
        if k <= 0:
            topk_by_user[user] = []
            continue
        top_items = torch.topk(user_scores, k=k, dim=0).indices.tolist()
        topk_by_user[user] = [int(item) for item in top_items]

    return topk_by_user