predictive_coding.py

import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Predictive Coding Layer
class PCLayer(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super(PCLayer, self).__init__()
        self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
        if self.bias is not None:
            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / (fan_in**0.5)
            nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, x):
        # Generative prediction: mu_l = f(W_l * x_{l-1} + b_l)
        out = torch.mm(x, self.weight.t())
        if self.bias is not None:
            out += self.bias
        return out

class PCNetwork(nn.Module):
    def __init__(self, layers_dims):
        super(PCNetwork, self).__init__()
        self.layers = nn.ModuleList()
        for i in range(len(layers_dims) - 1):
            self.layers.append(PCLayer(layers_dims[i], layers_dims[i+1]))
        
        self.activation = nn.ReLU()

    def forward(self, x):
        # Standard forward pass to initialize hidden states
        activations = [x]
        for layer in self.layers:
            x = self.activation(layer(x))
            activations.append(x)
        return activations

def main():
    # Parameters
    batch_size = 64
    epochs = 20
    lr_state = 0.1
    lr_weight = 0.0002
    inference_steps = 60
    
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")
    print(f"Using device: {device}")

    # Data - Neuromorphic systems often handle noise well
    train_transform = transforms.Compose([
        transforms.RandomRotation(15),
        transforms.RandomAffine(0, translate=(0.1, 0.1), scale=(0.9, 1.1)),
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=train_transform)
    test_dataset = datasets.MNIST('./data', train=False, transform=test_transform)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

    # Model: 784 -> 512 -> 512 -> 10
    model = PCNetwork([28 * 28, 512, 512, 10]).to(device)
    criterion = nn.CrossEntropyLoss()

    # Training loop
    for epoch in range(1, epochs + 1):
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            data = data.view(-1, 28 * 28).to(device)
            target = target.to(device)
            
            # One-hot encoding for targets (PC specific)
            target_onehot = torch.zeros(data.size(0), 10).to(device)
            target_onehot.scatter_(1, target.view(-1, 1), 1)
            
            # 1. Initialization: Standard forward pass
            with torch.no_grad():
                states = model(data) # [x0, x1, x2, ...]
            
            # Detach and make hidden states learnable for the inference phase
            states = [s.detach().clone().requires_grad_(True) for s in states]
            
            # 2. Inference Phase: Minimize Prediction Error
            # Clamping output (target_onehot)
            with torch.no_grad():
                states[-1].copy_(target_onehot)
            
            for _ in range(inference_steps):
                energy = 0
                for i in range(len(model.layers)):
                    mu = model.activation(model.layers[i](states[i]))
                    error = states[i+1] - mu
                    energy += torch.sum(error**2)
                
                # Gradient descent on hidden states
                energy.backward()
                with torch.no_grad():
                    for i in range(1, len(states) - 1): # Don't update clamped input/output
                        states[i] -= lr_state * states[i].grad
                        states[i].grad.zero_()
            
            # 3. Learning Phase: Update Weights
            energy = 0
            for i in range(len(model.layers)):
                mu = model.activation(model.layers[i](states[i]))
                error = states[i+1] - mu
                energy += torch.sum(error**2)
            
            model.zero_grad()
            energy.backward()
            
            with torch.no_grad():
                for layer in model.layers:
                    layer.weight -= lr_weight * layer.weight.grad
                    if layer.bias is not None:
                        layer.bias -= lr_weight * layer.bias.grad
                    layer.weight.grad.zero_()
                    if layer.bias is not None:
                        layer.bias.grad.zero_()

            if batch_idx % 200 == 0:
                print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                      f'({100. * batch_idx / len(train_loader):.0f}%)]\tEnergy: {energy.item():.6f}')

        # Testing loop
        model.eval()
        test_loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in test_loader:
                data = data.view(-1, 28 * 28).to(device)
                target = target.to(device)
                states = model(data)
                output = states[-1]
                test_loss += criterion(output, target).item()
                pred = output.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()

        test_loss /= len(test_loader)
        print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} '
              f'({100. * correct / len(test_loader.dataset):.2f}%)\n')

if __name__ == "__main__":
    main()