Use cross features in a complete predictor

Feature interactions

Predict a click using interactions between person, item, and context features. Input x [B, D] contains B examples with D numeric features each. A cross branch repeatedly scales the original features by a learned weighted sum of the current features, then adds a bias and the previous result. A separate deep branch transforms the same input with a linear layer and ReLU. Concatenate both outputs and use a head to return B logits: unrestricted click scores. The main section trains the complete model for one illustrative update.

View the full file · deep_cross_model.py

1import torch2from torch import nn3 4class CrossNetworkV1(nn.Module):5    def __init__(self, dim, depth=2):6        super().__init__()7        self.weights = nn.ParameterList([8            nn.Parameter(torch.randn(dim) * dim ** -0.5) for _ in range(depth)9        ])10        self.biases = nn.ParameterList([11            nn.Parameter(torch.zeros(dim)) for _ in range(depth)12        ])13 14    def forward(self, x0):15        # x0: [B, D]; each row is one example's combined features.16        x = x017        for weight, bias in zip(self.weights, self.biases):18            projection = (x * weight).sum(dim=-1, keepdim=True)  # [B, 1]19            x = x0 * projection + bias + x                     # [B, D]20        return x21 22class DeepCrossPredictor(nn.Module):23    def __init__(self, dim):24        super().__init__()25        self.cross = CrossNetworkV1(dim, depth=2)26        self.deep = nn.Sequential(nn.Linear(dim, 16), nn.ReLU())27        self.head = nn.Linear(dim + 16, 1)28 29    def forward(self, x):30        combined = torch.cat([self.cross(x), self.deep(x)], dim=-1)31        return self.head(combined).squeeze(-1)  # [B]32 33if __name__ == "__main__":34    torch.manual_seed(0)35    x = torch.tensor([[2., 3.], [1., 4.]])  # [B=2, D=2]36    labels = torch.tensor([1., 0.])37    model = DeepCrossPredictor(dim=2)38    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)39    optimizer.zero_grad()40    logits = model(x)41    loss = nn.BCEWithLogitsLoss()(logits, labels)42    loss.backward()43    optimizer.step()44    with torch.no_grad():45        print("Probabilities:", model(x).sigmoid())46    # A toy update shows the wiring, not prediction quality.

Lines 22–27
Compose the two branches
The cross branch returns [B, D]. The deep branch maps [B, D] to [B, 16] using a learned Linear layer and ReLU, which keeps positive values and sets negative ones to zero. Concatenating along dim=-1 gives [B, D+16]. These are two complementary ways to transform the same features.
Two parallel branches, one prediction head
1. Send the same input to both branches
  Numeric featuresx [B, D]
2. Each branch transforms the input independently
  Cross layers[B, D]
  Linear + ReLU[B, 16]
3. Concatenate along the feature axis
  Combined representation[B, D + 16]
Concatenation keeps both representations side by side. Adding the branches would need equal widths and would mix their features together.
Python
```
20        return x21 22class DeepCrossPredictor(nn.Module):23    def __init__(self, dim):24        super().__init__()25        self.cross = CrossNetworkV1(dim, depth=2)26        self.deep = nn.Sequential(nn.Linear(dim, 16), nn.ReLU())27        self.head = nn.Linear(dim + 16, 1)28 29    def forward(self, x):
```
Quick checkWith D = 4 input features, how wide is combined?

deep_cross_model.py

1import torch2from torch import nn3 4class CrossNetworkV1(nn.Module):5    def __init__(self, dim, depth=2):6        super().__init__()7        self.weights = nn.ParameterList([8            nn.Parameter(torch.randn(dim) * dim ** -0.5) for _ in range(depth)9        ])10        self.biases = nn.ParameterList([11            nn.Parameter(torch.zeros(dim)) for _ in range(depth)12        ])13 14    def forward(self, x0):15        # x0: [B, D]; each row is one example's combined features.16        x = x017        for weight, bias in zip(self.weights, self.biases):18            projection = (x * weight).sum(dim=-1, keepdim=True)  # [B, 1]19            x = x0 * projection + bias + x                     # [B, D]20        return x21 22class DeepCrossPredictor(nn.Module):23    def __init__(self, dim):24        super().__init__()25        self.cross = CrossNetworkV1(dim, depth=2)26        self.deep = nn.Sequential(nn.Linear(dim, 16), nn.ReLU())27        self.head = nn.Linear(dim + 16, 1)28 29    def forward(self, x):30        combined = torch.cat([self.cross(x), self.deep(x)], dim=-1)31        return self.head(combined).squeeze(-1)  # [B]32 33if __name__ == "__main__":34    torch.manual_seed(0)35    x = torch.tensor([[2., 3.], [1., 4.]])  # [B=2, D=2]36    labels = torch.tensor([1., 0.])37    model = DeepCrossPredictor(dim=2)38    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)39    optimizer.zero_grad()40    logits = model(x)41    loss = nn.BCEWithLogitsLoss()(logits, labels)42    loss.backward()43    optimizer.step()44    with torch.no_grad():45        print("Probabilities:", model(x).sigmoid())46    # A toy update shows the wiring, not prediction quality.

Compose the two branches