import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms class AgenticMLPClassifier(nn.Module): """ Agentic MLP for MNIST. A controller agent observes the input and outputs per-layer affine modulation parameters (gamma, beta), dynamically reconfiguring the base MLP's forward pass for each individual sample. """ def __init__(self, input_dim=784, hidden_dim=256, num_classes=10, num_layers=3): super().__init__() self.input_dim = input_dim self.hidden_dim = hidden_dim self.num_layers = num_layers # --- Base Feature MLP (the "environment" the agent acts upon) --- self.base_layers = nn.ModuleList() self.base_layers.append(nn.Linear(input_dim, hidden_dim)) for _ in range(num_layers - 1): self.base_layers.append(nn.Linear(hidden_dim, hidden_dim)) self.classifier = nn.Linear(hidden_dim, num_classes) # --- Agent Controller (observes image, plans transformations) --- # For each layer we output hidden_dim scales and hidden_dim shifts. agent_out_dim = num_layers * hidden_dim * 2 self.agent = nn.Sequential( nn.Linear(input_dim, 128), nn.ReLU(), nn.Linear(128, 128), nn.ReLU(), nn.Linear(128, agent_out_dim) ) self._init_weights() def _init_weights(self): for m in self.modules(): if isinstance(m, nn.Linear): nn.init.kaiming_normal_(m.weight, nonlinearity='relu', mode='fan_in') if m.bias is not None: nn.init.constant_(m.bias, 0.0) def forward(self, x): # x: (N, 1, 28, 28) x_flat = x.view(x.size(0), -1) # Agent perceives image and emits actions for every layer agent_out = self.agent(x_flat) # (N, num_layers*2*H) agent_out = agent_out.view(x.size(0), self.num_layers, 2, self.hidden_dim) # Residual modulation so base network can learn a stable prior gammas = 1.0 + torch.tanh(agent_out[:, :, 0, :]) # (N, L, H): explore around 1 betas = torch.tanh(agent_out[:, :, 1, :]) # (N, L, H): explore around 0 # Agentic forward pass through base MLP h = x_flat for i, layer in enumerate(self.base_layers): h = layer(h) h = gammas[:, i, :] * h + betas[:, i, :] # agentic modulation h = torch.relu(h) logits = self.classifier(h) return logits def train_epoch(model, loader, optimizer, criterion, device): model.train() total_loss, correct, total = 0.0, 0, 0 for data, target in loader: data, target = data.to(device), target.to(device) for i in range(1): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() if i==0: total_loss += loss.item() * data.size(0) _, pred = output.max(1) total += target.size(0) correct += pred.eq(target).sum().item() return total_loss / total, 100. * correct / total def train_epoch10(model, loader, optimizer, criterion, device): model.train() total_loss, correct, total = 0.0, 0, 0 for data, target in loader: data, target = data.to(device), target.to(device) for i in range(2): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() if i==0: total_loss += loss.item() * data.size(0) _, pred = output.max(1) total += target.size(0) correct += pred.eq(target).sum().item() return total_loss / total, 100. * correct / total def evaluate(model, loader, criterion, device): model.eval() total_loss, correct, total = 0.0, 0, 0 with torch.no_grad(): for data, target in loader: data, target = data.to(device), target.to(device) output = model(data) total_loss += criterion(output, target).item() * data.size(0) _, pred = output.max(1) total += target.size(0) correct += pred.eq(target).sum().item() return total_loss / total, 100. * correct / total def main(): # --- Config --- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') batch_size = 100 epochs = 10 lr = 1e-3 # --- Data --- transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_ds = datasets.MNIST('../data', train=True, download=True, transform=transform) test_ds = datasets.MNIST('../data', train=False, download=True, transform=transform) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=0) test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=0) # --- Model, Loss, Optimizer --- model = AgenticMLPClassifier(input_dim=784, hidden_dim=100, num_classes=10, num_layers=3) model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) print("=" * 60) print("Training Agentic MLP on MNIST") print(f"Device: {device} | Params: {sum(p.numel() for p in model.parameters()) / 1e3:.1f}K") print("=" * 60) # --- Train + Test Loop --- for epoch in range(1, epochs + 1): train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion, device) if train_acc>0.96: train_loss, train_acc = train_epoch10(model, train_loader, optimizer, criterion, device) test_loss, test_acc = evaluate(model, test_loader, criterion, device) scheduler.step() print(f"Epoch {epoch:02d}/{epochs} | " f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.2f}% | " f"Test Loss: {test_loss:.4f}, Acc: {test_acc:.2f}%") # --- Final Report --- _, final_acc = evaluate(model, test_loader, criterion, device) print("=" * 60) print(f"Final Test Accuracy: {final_acc:.2f}%") print("=" * 60) if __name__ == "__main__": main()