Loading your learning…
Loading your learning…
Move from supplied feature vectors to item IDs. TargetAwareEmbedding owns a shared embedding table, a relevance scorer, padding masks, and weighted pooling. ClickPredictor wraps it with the prediction head. Inputs are history IDs [B, N] and target IDs [B]; outputs are logits [B] and inspectable history weights [B, N]. This complete example includes initialization, one training update, and inference. B counts examples, N counts padded history positions, and D is the embedding width. A logit is an unrestricted score; sigmoid converts it to a click probability.
import torchfrom torch import nn class TargetAwareEmbedding(nn.Module): def __init__(self, num_items, dim, padding_idx=0): super().__init__() self.padding_idx = padding_idx self.embedding = nn.Embedding(num_items, dim, padding_idx=padding_idx) self.score = nn.Sequential( nn.Linear(4 * dim, 32), nn.PReLU(), nn.Linear(32, 1) ) def forward(self, history_ids, target_ids): # history_ids: [B, N]; target_ids: [B]. Targets must be real items. valid = history_ids.ne(self.padding_idx) # [B, N] items = self.embedding(history_ids) # [B, N, D] target = self.embedding(target_ids) # [B, D] query = target[:, None, :].expand_as(items) # [B, N, D] features = torch.cat( [items, query, items - query, items * query], dim=-1 ) # [B, N, 4D] weights = self.score(features).squeeze(-1) # [B, N] weights = weights.masked_fill(~valid, 0.0) # No softmax interest = (weights[..., None] * items).sum(dim=1) # [B, D] return interest, target, weights class ClickPredictor(nn.Module): def __init__(self, num_items, dim): super().__init__() self.history = TargetAwareEmbedding(num_items, dim) self.head = nn.Sequential( nn.Linear(2 * dim, 32), nn.ReLU(), nn.Linear(32, 1) ) def forward(self, history_ids, target_ids): interest, target, weights = self.history(history_ids, target_ids) features = torch.cat([interest, target], dim=-1) # [B, 2D] logits = self.head(features).squeeze(-1) # [B] return logits, weights if __name__ == "__main__": torch.manual_seed(0) # Two examples, each with one history and one candidate. ID 0 is padding. history_ids = torch.tensor([[1, 2, 0], [3, 4, 5]]) # [B=2, N=3] target_ids = torch.tensor([3, 2]) # [B=2] clicked = torch.tensor([1., 0.]) # Observed labels model = ClickPredictor(num_items=6, dim=8) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) model.train() optimizer.zero_grad() logits, weights = model(history_ids, target_ids) loss = nn.BCEWithLogitsLoss()(logits, clicked) loss.backward() optimizer.step() model.eval() with torch.no_grad(): logits, weights = model(history_ids, target_ids) print("History weights:", weights) # [2, 3] print("Click probabilities:", logits.sigmoid()) # [2] # One toy update demonstrates wiring, not a useful trained recommender.history_ids [B, N] contains B examples with up to N history items each. target_ids [B] gives one candidate per example; clicked [B] records whether that candidate was clicked. N is the padded history length, not the number of predictions. Here B = 2: the first history is [1, 2] followed by padding 0, and the second is [3, 4, 5]. Their candidates are 3 and 2. All IDs use the same vocabulary, with 0 reserved for padding. Use integer tensors; targets must be non-padding IDs. The model returns two click logits, one per example.
if __name__ == "__main__": torch.manual_seed(0) # Two examples, each with one history and one candidate. ID 0 is padding. history_ids = torch.tensor([[1, 2, 0], [3, 4, 5]]) # [B=2, N=3] target_ids = torch.tensor([3, 2]) # [B=2] clicked = torch.tensor([1., 0.]) # Observed labels model = ClickPredictor(num_items=6, dim=8) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) model.train()TargetAwareEmbedding groups two trainable parts: an item embedding table and a network that scores each history item against a candidate. The table maps an item ID to D learned numbers. Using the same table for history and target IDs puts both in the same feature space. Unlike a plain EmbeddingBag sum or mean, this summary depends on the target. EmbeddingBag can accept supplied per-item weights in sum mode; it does not learn those weights from a target itself. Here the module performs lookup, learns the weights, and pools the result. It materializes item vectors and comparison features, so it does not have the fused memory efficiency of EmbeddingBag.
from torch import nn class TargetAwareEmbedding(nn.Module): def __init__(self, num_items, dim, padding_idx=0): super().__init__() self.padding_idx = padding_idx self.embedding = nn.Embedding(num_items, dim, padding_idx=padding_idx) self.score = nn.Sequential( nn.Linear(4 * dim, 32), nn.PReLU(), nn.Linear(32, 1) ) def forward(self, history_ids, target_ids):Histories can have different lengths. We fill unused positions with a reserved padding ID so they fit into one rectangular tensor history_ids [B, N]. Here that padding ID is 0. .ne means “not equal”: history_ids.ne(self.padding_idx) compares every ID with the padding ID, just like history_ids != self.padding_idx. The result, valid, is a Boolean tensor of shape [B, N]: True for a real history item and False for padding. For history_ids = [[1, 2, 0], [3, 0, 0]], valid is [[True, True, False], [True, False, False]]. It marks which positions may contribute; it does not remove positions or score their relevance. Later, we use this mask to zero out padding weights before pooling.
) def forward(self, history_ids, target_ids): # history_ids: [B, N]; target_ids: [B]. Targets must be real items. valid = history_ids.ne(self.padding_idx) # [B, N] items = self.embedding(history_ids) # [B, N, D] target = self.embedding(target_ids) # [B, D]