class FeatureEncoder(nn.Module):
'''Standardises each feature using the training rows, then embeds each scalar cell with a linear layer.'''
def __init__(self, embedding_size):
super().__init__()
self.linear_layer = nn.Linear(1, embedding_size)
def forward(self, x, train_test_split_index):
x = x.unsqueeze(-1) # (B, R, C, 1)
mean = x[:, :train_test_split_index].mean(dim=1, keepdim=True)
std = x[:, :train_test_split_index].std(dim=1, keepdim=True) + 1e-8
x = torch.clip((x - mean) / std, -100, 100)
return self.linear_layer(x) # (B, R, C, E)
class TargetEncoder(nn.Module):
'''Pads the unknown targets of the test rows with the training mean, then embeds every target cell.'''
def __init__(self, embedding_size):
super().__init__()
self.linear_layer = nn.Linear(1, embedding_size)
def forward(self, y_train, num_rows):
mean = y_train.mean(dim=1, keepdim=True)
padding = mean.repeat(1, num_rows - y_train.shape[1], 1)
y = torch.cat([y_train, padding], dim=1).unsqueeze(-1) # (B, R, 1, 1)
return self.linear_layer(y) # (B, R, 1, E)
class TransformerEncoderLayer(nn.Module):
'''Attention between features, attention between datapoints, then an MLP. Each with residual + LayerNorm.'''
def __init__(self, embedding_size, nhead, mlp_hidden_size):
super().__init__()
self.self_attention_between_datapoints = nn.MultiheadAttention(embedding_size, nhead, batch_first=True)
self.self_attention_between_features = nn.MultiheadAttention(embedding_size, nhead, batch_first=True)
self.linear1 = nn.Linear(embedding_size, mlp_hidden_size)
self.linear2 = nn.Linear(mlp_hidden_size, embedding_size)
self.norm1 = nn.LayerNorm(embedding_size)
self.norm2 = nn.LayerNorm(embedding_size)
self.norm3 = nn.LayerNorm(embedding_size)
def forward(self, src, train_test_split_index):
B, R, C, E = src.shape
# --- attention between features: every row is a sequence of C cells
src = src.reshape(B * R, C, E)
src = self.self_attention_between_features(src, src, src)[0] + src
src = self.norm1(src.reshape(B, R, C, E))
# --- attention between datapoints: every column is a sequence of R cells
src = src.transpose(1, 2).reshape(B * C, R, E)
train, test = src[:, :train_test_split_index], src[:, train_test_split_index:]
train_out = self.self_attention_between_datapoints(train, train, train)[0] # train rows attend to train rows
test_out = self.self_attention_between_datapoints(test, train, train)[0] # test rows attend to train rows only
src = torch.cat([train_out, test_out], dim=1) + src
src = self.norm2(src.reshape(B, C, R, E).transpose(1, 2))
# --- position-wise MLP
src = self.linear2(F.gelu(self.linear1(src))) + src
return self.norm3(src)
class Decoder(nn.Module):
def __init__(self, embedding_size, mlp_hidden_size, num_outputs):
super().__init__()
self.linear1 = nn.Linear(embedding_size, mlp_hidden_size)
self.linear2 = nn.Linear(mlp_hidden_size, num_outputs)
def forward(self, x):
return self.linear2(F.gelu(self.linear1(x)))
class NanoTabPFNModel(nn.Module):
def __init__(self, embedding_size, num_attention_heads, mlp_hidden_size, num_layers, num_outputs):
super().__init__()
self.feature_encoder = FeatureEncoder(embedding_size)
self.target_encoder = TargetEncoder(embedding_size)
self.transformer_blocks = nn.ModuleList(
[TransformerEncoderLayer(embedding_size, num_attention_heads, mlp_hidden_size) for _ in range(num_layers)]
)
self.decoder = Decoder(embedding_size, mlp_hidden_size, num_outputs)
def forward(self, x, y_train, train_test_split_index):
'''
x: (B, R, C) all rows, training rows first
y_train: (B, train_test_split_index) labels of the training rows
returns: (B, R - train_test_split_index, num_outputs) logits for the test rows
'''
if y_train.dim() == 2:
y_train = y_train.unsqueeze(-1)
x_emb = self.feature_encoder(x, train_test_split_index) # (B, R, C, E)
y_emb = self.target_encoder(y_train, x.shape[1]) # (B, R, 1, E)
src = torch.cat([x_emb, y_emb], dim=2) # (B, R, C+1, E)
for block in self.transformer_blocks:
src = block(src, train_test_split_index)
return self.decoder(src[:, train_test_split_index:, -1, :]) # target cells of the test rows
def make_model(num_outputs=MAX_CLASSES):
return NanoTabPFNModel(embedding_size=192, num_attention_heads=6, mlp_hidden_size=768, num_layers=6, num_outputs=num_outputs)
model = make_model().to(device)
print(f"{sum(p.numel() for p in model.parameters()) / 1e6:.2f} M parameters")
