import torch import torch.nn as nn import torch.optim as optim import numpy as np from torch.utils.data import DataLoader, TensorDataset # ------------------------------ # 1. Toy data with variable feature sizes # ------------------------------ vocab_size = 20 # words 0..19 n = 3 # lag # Generate random sentences num_sentences = 200 sentence_len = 10 sentences = [] for _ in range(num_sentences): sent = np.random.randint(0, vocab_size, sentence_len).tolist() sentences.append(sent) # Assign plateaus: each word gets a deterministic height but we will use # different numbers of plateau levels during training (2, 4, 8 levels) def get_height(word, num_levels): # Map word index to a discrete plateau level (0..num_levels-1) level = word % num_levels # Convert to a scalar in [0,1] return level / (num_levels - 1) if num_levels > 1 else 0.5 # Build sliding windows for a given number of plateau levels def build_dataset(sentences, n, num_levels): X_words, X_heights, y_words = [], [], [] for sent in sentences: for i in range(len(sent) - n): in_words = sent[i:i+n] heights = [get_height(w, num_levels) for w in in_words] target = sent[i+n] X_words.append(in_words) X_heights.append(heights) y_words.append(target) return (torch.tensor(X_words, dtype=torch.long), torch.tensor(X_heights, dtype=torch.float), torch.tensor(y_words, dtype=torch.long)) # We will create three datasets with different feature sizes (num_levels) datasets = { 'small': build_dataset(sentences, n, num_levels=2), # binary plateau 'medium': build_dataset(sentences, n, num_levels=4), # 4 levels 'large': build_dataset(sentences, n, num_levels=8) # 8 levels } # ------------------------------ # 2. Model that handles variable input feature size # ------------------------------ class VariableFeatureLSTM(nn.Module): def __init__(self, vocab_size, embed_dim=8, hidden_dim=16, max_feature_dim=2): """ max_feature_dim: maximum number of height scalars we might receive. But here we keep height as a scalar; instead we vary the number of plateau levels, which changes the distribution but not the dimension. To truly vary feature size, we can concatenate additional random features. """ super().__init__() self.word_embed = nn.Embedding(vocab_size, embed_dim) # We will always have height (1 dim) + optionally extra synthetic features. # Keep the configured width consistent with the projection layer. self.max_extra = max_feature_dim # Projection layer to map variable input to fixed hidden size self.input_proj = nn.Linear(embed_dim + 1 + self.max_extra, hidden_dim) self.lstm = nn.LSTM(hidden_dim, hidden_dim, batch_first=True) self.fc = nn.Linear(hidden_dim, vocab_size) def forward(self, word_seq, height_seq, extra_features=None): # word_seq: (batch, n) # height_seq: (batch, n, 1) normally batch, n = word_seq.shape word_emb = self.word_embed(word_seq) # (batch, n, embed_dim) height_seq = height_seq.unsqueeze(-1) # (batch, n, 1) if extra_features is None: extra_features = torch.zeros( batch, n, self.max_extra, device=word_seq.device, dtype=word_emb.dtype, ) else: extra_features = extra_features.to( device=word_seq.device, dtype=word_emb.dtype, ) # extra_features: (batch, n, extra_dim) extra_dim = extra_features.shape[-1] if extra_dim < self.max_extra: pad = torch.zeros( batch, n, self.max_extra - extra_dim, device=extra_features.device, dtype=extra_features.dtype, ) extra_features = torch.cat([extra_features, pad], dim=-1) elif extra_dim > self.max_extra: extra_features = extra_features[..., :self.max_extra] combined = [word_emb, height_seq, extra_features] x = torch.cat(combined, dim=-1) # (batch, n, embed_dim+1+extra) # Project to fixed hidden size x = self.input_proj(x) # (batch, n, hidden_dim) lstm_out, _ = self.lstm(x) # (batch, n, hidden_dim) last_out = lstm_out[:, -1, :] logits = self.fc(last_out) return logits # ------------------------------ # 3. Training with mixed feature sizes # ------------------------------ model = VariableFeatureLSTM(vocab_size) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.005) # Create a combined dataloader that samples from all three plateau levels # and optionally adds extra random features to further vary feature size def mixed_dataloader(batch_size=16): # Each item: (words, heights, targets, size_tag) items = [] for size_tag, (Xw, Xh, y) in datasets.items(): for i in range(len(Xw)): # Decide randomly to add extra features (0,1,2 dimensions) extra_dim = np.random.choice([0,1,2]) extra = None if extra_dim > 0: extra = torch.randn(1, n, extra_dim) # random noise items.append((Xw[i], Xh[i], y[i], extra, size_tag)) # Shuffle and batch np.random.shuffle(items) for i in range(0, len(items), batch_size): batch = items[i:i+batch_size] words = torch.stack([b[0] for b in batch]) heights = torch.stack([b[1] for b in batch]) targets = torch.stack([b[2] for b in batch]) extras = [b[3] for b in batch] # Padding for extra features (since dimensions may differ) max_extra = max([e.shape[-1] if e is not None else 0 for e in extras]) if max_extra > 0: padded_extras = [] for e in extras: if e is None: e = torch.zeros(1, n, max_extra) else: # pad to max_extra pad = torch.zeros(1, n, max_extra - e.shape[-1]) e = torch.cat([e, pad], dim=-1) padded_extras.append(e) extras_tensor = torch.cat(padded_extras, dim=0) else: extras_tensor = None yield words, heights, targets, extras_tensor # Training loop epochs = 300 for epoch in range(epochs): total_loss = 0.0 num_batches = 0 for words, heights, targets, extras in mixed_dataloader(batch_size=8): optimizer.zero_grad() logits = model(words, heights, extras) loss = criterion(logits, targets) loss.backward() optimizer.step() total_loss += loss.item() num_batches += 1 if (epoch+1) % 50 == 0: print(f"Epoch {epoch+1:3d}, Loss: {total_loss/num_batches:.4f}") # ------------------------------ # 4. Test generalization on unseen sentences # ------------------------------ test_sentences = [ [5, 12, 7, 19, 3, 8, 0, 15, 2, 11], [9, 1, 18, 4, 13, 6, 17, 10, 14, 16] ] def test_model(model, sentences, n, num_levels): correct = 0 total = 0 for sent in sentences: for i in range(len(sent)-n): in_words = sent[i:i+n] target = sent[i+n] heights = [get_height(w, num_levels) for w in in_words] words_t = torch.tensor([in_words], dtype=torch.long) heights_t = torch.tensor([heights], dtype=torch.float) with torch.no_grad(): # Test without extra features (use 0 extra dims) logits = model(words_t, heights_t, None) pred = torch.argmax(logits, dim=1).item() if pred == target: correct += 1 total += 1 return correct / total # Test on all three plateau levels for levels in [2,4,8]: acc = test_model(model, test_sentences, n, levels) print(f"Test accuracy for {levels} plateau levels: {acc:.3f}") # Also test with extra features during inference (to see if model can adapt) print("\nTesting with extra features (dim=2) at inference time:") for levels in [2,4,8]: # We need to create dummy extra features for the test sentences correct = 0 total = 0 for sent in test_sentences: for i in range(len(sent)-n): in_words = sent[i:i+n] target = sent[i+n] heights = [get_height(w, levels) for w in in_words] words_t = torch.tensor([in_words], dtype=torch.long) heights_t = torch.tensor([heights], dtype=torch.float) extra = torch.randn(1, n, 2) # random extra features with torch.no_grad(): logits = model(words_t, heights_t, extra) pred = torch.argmax(logits, dim=1).item() if pred == target: correct += 1 total += 1 print(f" {levels} levels: {correct/total:.3f}")