Problem with conditioning transformer

specialized_boy: <p>I implemented a transformer decoder and initially trained it on text sequences, achieving satisfactory performance. Subsequently, I aimed to condition the model on floating-point values. To accomplish this, I converted the floating-point values into embeddings using linear layers, which were then passed to the decoder as memory inputs. However, upon training with this setup, I observed an unusually low loss from the outset, leading me to suspect overfitting.</p>
<p>Despite my concerns, when I sampled from the model by providing a floating-point value and a start token to the decoder, the output sequences were unexpectedly repetitive, consisting of merely two alternating letters. This outcome suggested a problem beyond simple overfitting, as overfitting would typically result in the model reproducing identical sequences rather than generating repetitive patterns.</p>
<p>To diagnose the issue, I sampled outputs every 100 batches during training, and these sequences appeared to be normal, indicating that the training process itself may not be fundamentally flawed.</p>
<p>Given these observations, what do you think could be the problem with my approach?</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        
        self.dropout = nn.Dropout(p=dropout)
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0)]
        return self.dropout(x)

class ConditionalTransformerModel(nn.Module):
    def __init__(self, vocab_size, n_embd, n_head, n_layers, max_length, property_dim, property_hidden_dim, dropout=0.1):
        super().__init__()

        self.embed = nn.Embedding(vocab_size, n_embd)
        self.n_embd = n_embd
        self.property_processor = nn.Sequential(
            nn.Linear(property_dim, property_hidden_dim),
            nn.ReLU(),
            nn.Linear(property_hidden_dim, n_embd)
        )

        decoder_layer = nn.TransformerDecoderLayer(d_model=n_embd, nhead=n_head, dropout=dropout)
        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=n_layers)
        self.pos_encoder = PositionalEncoding(n_embd, dropout, max_length)
        self.generator = nn.Linear(n_embd, vocab_size)
        self.max_length = max_length

    def generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def forward(self, src, tgt, property):

        batch_size, seq_length = src.size(0), src.size(1)
        src = self.embed(src) * math.sqrt(self.n_embd)
        tgt = self.embed(tgt) * math.sqrt(self.n_embd)
        src = self.pos_encoder(src)
        tgt = self.pos_encoder(tgt)
        property_encoded = self.property_processor(property)
        property_encoded = property_encoded.unsqueeze(1).repeat(1, seq_length, 1)
        memory = property_encoded.transpose(0, 1)

        tgt_mask = self.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)

        output = self.transformer_decoder(tgt.transpose(0, 1), memory, tgt_mask=tgt_mask)
        output = output.transpose(0, 1)
        return self.generator(output)


# Define your model
model = ConditionalTransformerModel(vocab_size=vocab.vocab_size, n_embd=64, n_head=4, n_layers=4, max_length=512, property_dim=1, property_hidden_dim=128).to('cuda')

# Training loop
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

def sample_predictions(output, vocab, k=5):
    """Sample some predictions from the model output"""
    # Get the most probable next tokens (top-k)
    _, topk_indices = torch.topk(output, k, dim=-1)
    sampled_indices = topk_indices[:, :, 0]  # For simplicity, use the top-1 prediction
    sampled_tokens = [vocab.decode(indices.tolist()) for indices in sampled_indices]
    return sampled_tokens

for epoch in range(10):  # Let's still assume 10 epochs for simplicity
    model.train()  # Set the model to training mode
    total_loss = 0  # To accumulate loss over the epoch
    
    for i, (src_batch, tgt_batch, properties_batch) in enumerate(data_loader):
        src_batch, tgt_batch, properties_batch = src_batch.to('cuda'), tgt_batch.to('cuda'), properties_batch.to('cuda')
        
        optimizer.zero_grad()  # Zero the gradients
        
        # Adjust the model's forward pass call to match the updated signature.
        # Now src is used as input to the encoder and tgt is used as input to the decoder
        # Note: You need to ensure that your model's forward method accepts these arguments correctly
        output = model(src_batch, tgt_batch, properties_batch)  # src_batch is now explicitly separated from tgt_batch
        
        # Calculate loss. Note that we need to adjust the targets to match the output dimensions.
        # Since output is likely of shape [batch_size, seq_len, vocab_size] and tgt_batch is [batch_size, seq_len],
        # we transpose output to [batch_size, vocab_size, seq_len] before passing it along with tgt_batch to the criterion.
        loss = criterion(output.transpose(1, 2), tgt_batch)  # Adjust loss calculation to use tgt_batch
        
        loss.backward()  # Backpropagate the loss
        optimizer.step()  # Update model parameters
        total_loss += loss.item()  # Accumulate the loss
    
        if i % 100 == 0:  # Sample predictions every 100 batches
            sampled_tokens = sample_predictions(output, vocab)
            print(f"Batch {i}: Sample Predictions: {sampled_tokens[:5]}")  # Print sample predictions from the first 5 sequences

    # Print average loss for the epoch
    print(f"Epoch {epoch+1}, Loss: {total_loss / len(data_loader)}")
</code></pre>
<p>this is the code I am using any comments?</p>