37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
import torch
|
||
from torch import nn
|
||
|
||
# 2) Number "embedding" network: R -> R^d
|
||
class NumberEmbedder(nn.Module):
|
||
def __init__(self, d=4, hidden=16):
|
||
super().__init__()
|
||
self.net = nn.Sequential(
|
||
nn.Linear(1, hidden),
|
||
nn.ReLU(),
|
||
nn.Linear(hidden, d),
|
||
)
|
||
|
||
def forward(self, x):
|
||
return self.net(x)
|
||
|
||
# 3) Comparator head: takes (ea, eb, e) -> logit for "a > b"
|
||
class PairwiseComparator(nn.Module):
|
||
def __init__(self, d=4, hidden=16, k=1.0):
|
||
super().__init__()
|
||
self.log_k = nn.Parameter(torch.tensor([k]))
|
||
self.embed = NumberEmbedder(d, hidden)
|
||
self.head = nn.Sequential(
|
||
nn.Linear(d, hidden),
|
||
nn.ReLU(),
|
||
nn.Linear(hidden, hidden),
|
||
nn.ReLU(),
|
||
nn.Linear(hidden, 1),
|
||
)
|
||
|
||
def forward(self, a, b):
|
||
# trying to force antisym here: h(a,b)=−h(b,a)
|
||
phi = self.head(self.embed(a-b))
|
||
phi_neg = self.head(self.embed(b-a))
|
||
logit = phi - phi_neg
|
||
|
||
return (self.log_k ** 2) * logit |