[Python] Why my Transformer model did not work well when dealing with single cell multi-omic...

The complete codes and data are available at:Google Disk

I'm working on a high-dimensional regression problem and have built a Transformer-based model in PyTorch. While the model trains, I'm observing a strange behavior: the validation loss finds a reasonable value in the first few epochs and then gets completely stuck for the rest of the training, not changing at all.

For example, the log for one fold looks like this:

===== Fold 2 ===== Epoch 10/30, Val Loss: -0.8088 Epoch 20/30, Val Loss: -0.8088 Epoch 30/30, Val Loss: -0.8088 This suggests the model isn't learning anything new that generalizes to the validation set after an initial phase. I'm trying to understand the potential causes.

My Question:

What are the most likely reasons for the validation loss to plateau so perfectly and early?

Is it a sign of rapid overfitting? If so, what are the best strategies to combat this in a Transformer for tabular data (beyond the dropout I'm already using)?

Could this be an issue with the custom loss function? Since Pearson correlation is scale-invariant, could the model be finding a "good enough" pattern very quickly, after which the gradients become too small to allow for further learning?

Is there a known issue or a better technique for using Transformers on extremely high-dimensional, sparse data like this? For example, is the simple linear projection at the start a potential bottleneck?

Any insights or debugging suggestions would be greatly appreciated.

Problem Background:

Task: This is for the Open Problems in Single-Cell Analysis Kaggle competition. The goal is to predict gene expression levels (~23k features) from ATAC-seq data (~150k features). The data is very large and sparse.

Algorithm: I'm using a 5-fold cross-validation setup. The model is trained using a custom loss function designed to maximize the Pearson correlation coefficient.

Loss Function: The loss is -pearson_correlation. So, a validation loss of -0.8088 means the model has achieved a strong Pearson correlation of 0.8088, but it simply stops improving beyond that point.

Model and Training Code:

Here is the core architecture of the TabularTransformer, the custom loss function, and the training loop.

1. The Transformer Model The model projects the high-dimensional input vector into a sequence, adds a [CLS] token and positional encodings, and feeds it through a standard Transformer Encoder.

Python

import torch
import torch.nn as nn

class TabularTransformer(nn.Module):
def __init__(self, num_features, num_targets, seq_len=16, d_model=256, nhead=8, num_layers=3, dim_feedforward=512, dropout=0.1):
super().__init__()
self.projector = nn.Linear(num_features, seq_len * d_model)
self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
self.pos_encoder = nn.Parameter(torch.randn(1, seq_len + 1, d_model))

encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, batch_first=True)
self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

self.mlp_head = nn.Sequential(
nn.LayerNorm(d_model),
nn.Linear(d_model, d_model // 2),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(d_model // 2, num_targets)
)
self.d_model = d_model
self.seq_len = seq_len

def forward(self, x):
x = self.projector(x)
x = x.reshape(-1, self.seq_len, self.d_model)

batch_size = x.size(0)
cls_tokens = self.cls_token.expand(batch_size, -1, -1)
x = torch.cat((cls_tokens, x), dim=1)
x += self.pos_encoder
x = self.transformer_encoder(x)
cls_output = x[:, 0]
output = self.mlp_head(cls_output)
return output

1. Custom Loss Function This loss calculates Pearson correlation by mean-centering the vectors and then applying cosine similarity. The negative is taken so it can be minimized.

Python

def cosine_similarity_loss(y_true, y_pred):
y_true_centered = y_true - torch.mean(y_true, dim=1, keepdim=True)
y_pred_centered = y_pred - torch.mean(y_pred, dim=1, keepdim=True)

y_true_norm = torch.nn.functional.normalize(y_true_centered, p=2, dim=1)
y_pred_norm = torch.nn.functional.normalize(y_pred_centered, p=2, dim=1)

return -torch.nn.CosineSimilarity(dim=1)(y_true_norm, y_pred_norm).mean()

1. The Training Loop Snippet (within one fold) This is a standard training loop with a ReduceLROnPlateau scheduler.

Python

(Inside a K-Fold loop)


model = TabularTransformer(**model_params).to(DEVICE)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=6)

for epoch in range(epochs):
model.train()
# ... Training loop over batches, loss.backward(), optimizer.step() ...

model.eval()
val_loss = 0
with torch.no_grad():
for features, targets in valid_loader:
features, targets = features.to(DEVICE), targets.to(DEVICE)
predictions = model(features)
loss = loss_fn(targets, predictions)
val_loss += loss.item() * len(targets)

val_loss /= len(valid_dataset)
scheduler.step(val_loss) # Scheduler should trigger if loss plateaus

print(f"Epoch {epoch+1}, Val Loss: {val_loss:.4f}")

Continue reading...