Loading your learning…
Loading your learning…
Predict a click from item IDs and known categories. B is the number of examples, N the padded history length, D the learned embedding width, and G the category count. Supply history IDs [B, N], one target ID per example [B], and category IDs [B, N]. The model looks up D-feature vectors, averages relevant items within categories, then weights and sums those summaries for the candidate. A head returns logits [B]: unrestricted click scores. Observed binary labels [B] train the whole model with one loss.
import torchfrom torch import nn class Relevance(nn.Module): def __init__(self, dim): super().__init__() self.net = nn.Sequential(nn.Linear(4 * dim, 16), nn.ReLU(), nn.Linear(16, 1)) def forward(self, items, target): query = target[:, None, :].expand_as(items) features = torch.cat([items, query, items - query, items * query], dim=-1) return self.net(features).squeeze(-1).sigmoid() class HierarchicalInterest(nn.Module): def __init__(self, dim, categories): super().__init__() self.categories = categories self.item_score = Relevance(dim) self.category_score = Relevance(dim) def forward(self, items, target, category_ids): # items [B, N, D], target [B, D], category_ids [B, N]; -1 means padding. groups = torch.arange(self.categories, device=items.device) membership = category_ids[:, None, :] == groups[None, :, None] # [B, G, N] item_weights = self.item_score(items, target) # [B, N] group_weights = item_weights[:, None, :] * membership # [B, G, N] mass = group_weights.sum(dim=-1, keepdim=True) # [B, G, 1] numerator = (group_weights[..., None] * items[:, None, :, :]).sum(dim=2) summaries = numerator / mass.clamp_min(torch.finfo(items.dtype).tiny) category_weights = self.category_score(summaries, target) # [B, G] category_weights = category_weights.masked_fill(mass.squeeze(-1) == 0, 0) interest = (category_weights[..., None] * summaries).sum(dim=1) # [B, D] return interest, group_weights, category_weights class HierarchicalRecommender(nn.Module): def __init__(self, num_items, dim, categories): super().__init__() self.embedding = nn.Embedding(num_items, dim, padding_idx=0) self.pool = HierarchicalInterest(dim, categories) self.head = nn.Linear(2 * dim, 1) def forward(self, history_ids, target_ids, category_ids): # Valid category IDs: 0..G-1. Reserve item ID 0 for padding. category_ids = category_ids.masked_fill(history_ids == 0, -1) items, target = self.embedding(history_ids), self.embedding(target_ids) interest, _, _ = self.pool(items, target, category_ids) return self.head(torch.cat([interest, target], dim=-1)).squeeze(-1) if __name__ == "__main__": torch.manual_seed(0) histories = torch.tensor([[1, 2, 3], [4, 0, 0]]) # [B=2, N=3] targets = torch.tensor([4, 2]) # [B] categories = torch.tensor([[0, 0, 1], [1, -1, -1]]) labels = torch.tensor([1., 0.]) model = HierarchicalRecommender(num_items=5, dim=8, categories=2) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) optimizer.zero_grad() logits = model(histories, targets, categories) # [B] loss = nn.BCEWithLogitsLoss()(logits, labels) loss.backward() optimizer.step() with torch.no_grad(): print("Probabilities:", model(histories, targets, categories).sigmoid()) # One toy update is not a useful trained recommender.The item embedding table is shared by history and target IDs. The pool returns one interest [B, D], which is joined with the target [B, D]. The final Linear maps [B, 2D] to [B, 1]; squeeze(-1) gives B independent click logits while preserving a batch of size one.
The target affects relevance and also reaches the prediction head directly. A zero interest vector therefore does not force a zero click probability.
return interest, group_weights, category_weights class HierarchicalRecommender(nn.Module): def __init__(self, num_items, dim, categories): super().__init__() self.embedding = nn.Embedding(num_items, dim, padding_idx=0) self.pool = HierarchicalInterest(dim, categories) self.head = nn.Linear(2 * dim, 1) def forward(self, history_ids, target_ids, category_ids):Item ID 0 is padding, so the wrapper forces its category to −1 before pooling. Targets must be real item IDs. Nonpadding history items need valid category IDs from 0 to G−1; missing or out-of-range category values would exclude those items, so validate that contract in data preparation.
def forward(self, history_ids, target_ids, category_ids): # Valid category IDs: 0..G-1. Reserve item ID 0 for padding. category_ids = category_ids.masked_fill(history_ids == 0, -1) items, target = self.embedding(history_ids), self.embedding(target_ids) interest, _, _ = self.pool(items, target, category_ids) return self.head(torch.cat([interest, target], dim=-1)).squeeze(-1) if __name__ == "__main__":