-
Notifications
You must be signed in to change notification settings - Fork 1
/
RFM.py
159 lines (131 loc) · 6.71 KB
/
RFM.py
1
2
3
4
5
6
7
8
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
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
import torch
import torch.nn as nn
from torch.nn.init import xavier_normal_, constant_
from recbole.model.abstract_recommender import ContextRecommender
from recbole.model.loss import RegLoss
from torch.nn import functional as F
def RotationBasedAttention(theta_a, theta_b, weight=None):
cos_a, cos_b = torch.cos(theta_a), torch.cos(theta_b)
sin_a, sin_b = torch.sin(theta_a), torch.sin(theta_b)
if weight is not None:
cos_a, sin_a = cos_a * weight, sin_a * weight
return torch.sigmoid(cos_a @ cos_b.transpose(-2, -1) + sin_a @ sin_b.transpose(-2, -1)) # cos a-b = cosa cosb + sin a sin b
class SelfAttentiveRotation(nn.Module):
def __init__(self, config, input_shape, output_shape) -> None:
super().__init__()
self.input_shape, self.output_shape = input_shape, output_shape
field_num = self.field_num = config.num_field
self.head_num = config.head_num
self.initialize_parameters(field_num, input_shape, output_shape, config.drop_rate_att)
def initialize_parameters(self, field_num, input_shape, output_shape, dropout_prob):
self.Q, self.K, self.V = [
nn.Parameter(torch.ones(field_num, input_shape, output_shape))
for _ in range(3)
]
self.linear_transform = nn.Linear(input_shape, output_shape, bias=False)
self.dropout = nn.Dropout(p=dropout_prob)
self.norm = nn.LayerNorm([self.output_shape])
self.weight = nn.Parameter(torch.ones(1, self.output_shape // self.head_num))
for param in [self.Q, self.K, self.V, self.linear_transform.weight, self.weight]:
xavier_normal_(param, gain=1.414)
def forward(self, theta):
tensor = theta.reshape(theta.shape[0], -1, self.input_shape)
queries, keys, values = [torch.einsum('bfd,fdw->bfw', tensor, param) for param in [self.Q, self.K, self.V]]
head_tensors = [
torch.stack(torch.split(tensor, [self.output_shape // self.head_num] * self.head_num, dim=-1), dim=1)
for tensor in [queries, keys, values]
]
queries, keys, values = head_tensors
scores = RotationBasedAttention(queries, keys, weight=self.weight).transpose(-2, -1)
values = self.dropout(values).transpose(-2, -1)
output = torch.cat(
torch.split(
torch.transpose(values @ scores, -2, -1),
[1] * self.head_num, dim=1
), dim=-1
).squeeze(1)
theta = self.dropout(theta)
representation = self.norm(output + self.linear_transform(theta))
return representation
class RFM(ContextRecommender):
def __init__(self, config, dataset) -> None:
super().__init__(config, dataset)
self.attention_architecture = [config.embedding_size] + [config.hidden_units] * config.att_layers
config.num_field = self.num_feature_field
self.saro_layers = nn.Sequential(
*(SelfAttentiveRotation(config, in_shape, out_shape)
for in_shape, out_shape in zip(self.attention_architecture[:-1], self.attention_architecture[1:]))
)
self.norm1 = nn.LayerNorm([self.attention_architecture[-1]])
self.norm2 = nn.LayerNorm([self.attention_architecture[-1]])
config.mlp_list = [self.attention_architecture[-1] * config.num_field] + config.mlp_list
self.ampnet = AmpNet(config)
self.add_first_order_residual = config.add_first_order_residual
self.projection = nn.Linear(config.embedding_size, self.attention_architecture[-1])
nn.init.normal_(self.projection.weight, mean=0, std=0.01)
self.apply(self._initialize_weights)
self.sigmoid = nn.Sigmoid()
self.loss = nn.BCEWithLogitsLoss()
self.reg_loss = RegLoss()
def _initialize_weights(self, module):
if isinstance(module, nn.Embedding):
nn.init.xavier_normal_(module.weight.data)
def forward(self, interaction):
angular_embeddings = self.concat_embed_input_fields(interaction)
theta = self.saro_layers(angular_embeddings)
r, p = torch.cos(theta), torch.sin(theta)
if self.add_first_order_residual:
r = self.norm1(r + torch.cos(self.projection(angular_embeddings)))
p = self.norm2(p + torch.sin(self.projection(angular_embeddings)))
logits = self.ampnet((r, p)).squeeze(-1)
return logits
def calculate_loss(self, interaction):
label = interaction[self.LABEL]
output = self.forward(interaction)
return self.loss(output, label)
def predict(self, interaction):
return self.sigmoid(self.forward(interaction))
class Amp(nn.Module):
def __init__(self, config, input_shape, output_shape, activation=None) -> None:
super().__init__()
self.input_shape = input_shape
self.output_shape = output_shape
self.group_size = config.group_hidden_dimension
# Initialize linear layer
self.linear = nn.Linear(input_shape, output_shape)
nn.init.normal_(self.linear.weight, mean=0, std=0.01)
# Initialize dropout layer
self.dropout = nn.Dropout(p=config.dp_rate_amp)
# Activation function
self.activation = activation
# Normalization layers, here we use LayerNorm to implement GroupNorm
self.norm1 = nn.LayerNorm([self.group_size])
self.norm2 = nn.LayerNorm([self.group_size])
def forward(self, kws):
real, imag = kws
real = self.dropout(real).view(real.size(0), -1)
imag = self.dropout(imag).view(imag.size(0), -1)
real, imag = self.linear(real), self.linear(imag)
if self.activation:
real, imag = self.activation(real), self.activation(imag)
real = real.view(real.size(0), -1, self.group_size)
imag = imag.view(imag.size(0), -1, self.group_size)
real, imag = self.norm1(real), self.norm2(imag)
return real, imag
class AmpNet(nn.Module):
def __init__(self, config) -> None:
super().__init__()
self.mlp_config_list = config.mlp_list
mlps = [Amp(config, in_shape, out_shape, nn.ReLU())
for in_shape, out_shape in zip(self.mlp_config_list[:-1], self.mlp_config_list[1:])]
self.block = nn.Sequential(*mlps)
# Regularization layer
self.regularization = nn.Linear(self.mlp_config_list[-1], 1)
nn.init.xavier_normal_(self.regularization.weight)
def forward(self, feature):
real, imag = feature
real, imag = self.block((real, imag))
real = real.view(real.size(0), -1)
imag = imag.view(imag.size(0), -1)
real, imag = self.regularization(real), self.regularization(imag)
return real + imag