""" Fourier Interference Layer for MNIST ====================================== Key idea: Instead of learning raw weights W[i,j], we learn Fourier *amplitudes* A[k] and *phases* φ[k] for a set of spatial frequencies k. The weight field is reconstructed as a sum of interfering waves: W(x, y) = Σ_k A[k] · cos(2π k·(x,y) + φ[k]) This is exactly a truncated Fourier series. Individual waves are perfectly smooth, but their *superposition* produces sharp edges, localised blobs, and fine detail — the same way a square wave emerges from summing sin harmonics. Advantages over SPDE: - Sharp features are cheap (just add high-frequency terms) - Every frequency is independent and interpretable - Naturally hierarchical: low-k = global structure, high-k = fine detail - Exact reconstruction at any resolution (truly continuous) Advantages over raw linear weights: - Far fewer parameters for the same expressive capacity - Structured inductive bias: features must be spatially coherent - Can visualise exactly which frequencies matter """ import torch import torch.nn as nn import torch.nn.functional as F import torchvision import torchvision.transforms as transforms import matplotlib.pyplot as plt import numpy as np torch.manual_seed(42) # ───────────────────────────────────────────────────────────────── # Core: Fourier Interference Weight Field # ───────────────────────────────────────────────────────────────── class FourierInterferenceField(nn.Module): """ Learns a 2D weight field as a truncated Fourier series. For an (H × W) output field we keep all frequency pairs (kx, ky) with |kx| <= max_freq and |ky| <= max_freq. Total learnable params per output neuron: 2 × (2·max_freq+1)² (amplitude + phase for each frequency pair) Compare to a raw dense weight: H × W params per neuron. With max_freq=7, H=W=28: 2×225 = 450 vs 784 — 43% fewer params, but the Fourier basis can still reconstruct sharp edges. """ def __init__(self, H, W, max_freq=7): super().__init__() self.H, self.W = H, W self.max_freq = max_freq # Enumerate all frequency pairs (kx, ky) freqs_1d = torch.arange(-max_freq, max_freq + 1) # 2F+1 values KX, KY = torch.meshgrid(freqs_1d, freqs_1d, indexing='ij') # Shape (n_freqs, 2) — one (kx,ky) pair per row self.register_buffer('freq_pairs', torch.stack([KX.flatten(), KY.flatten()], dim=1).float()) n_freqs = self.freq_pairs.shape[0] # (2F+1)² # Learnable: amplitude A and phase φ per frequency # Initialise amplitudes small, phases random self.amplitudes = nn.Parameter(torch.randn(n_freqs) * 0.02) self.phases = nn.Parameter(torch.rand(n_freqs) * 2 * torch.pi) # Fixed spatial grid in [0,1]² shape (H*W, 2) gy = torch.linspace(0, 1, H) gx = torch.linspace(0, 1, W) GY, GX = torch.meshgrid(gy, gx, indexing='ij') # (H*W, 2) self.register_buffer('spatial_grid', torch.stack([GY.flatten(), GX.flatten()], dim=1)) def forward(self): """ Returns weight field of shape (H, W). For each spatial position p and frequency k: contribution = A[k] · cos( 2π · k·p + φ[k] ) The sum over k is a matrix-vector product after broadcasting: phase_arg = 2π · (spatial_grid @ freq_pairs.T) + phases (H*W, n_freqs) field = (cos(phase_arg) * amplitudes).sum(dim=-1) (H*W,) """ # (H*W, n_freqs): 2π · k·p for all positions and frequencies phase_arg = 2 * torch.pi * (self.spatial_grid @ self.freq_pairs.T) # Add learned phase offset and amplitude, sum over frequencies field = (torch.cos(phase_arg + self.phases) * self.amplitudes).sum(dim=-1) return field.reshape(self.H, self.W) # ───────────────────────────────────────────────────────────────── # Fourier Interference Linear Layer # ───────────────────────────────────────────────────────────────── class FourierInterferenceLinear(nn.Module): """ A linear layer whose weight matrix is built from Fourier interference. Each output neuron i gets its own independent field F_i(x,y), parameterised by amplitudes A_i[k] and phases φ_i[k]. The weight matrix W has shape (out_features, in_features), where in_features corresponds to flattened 2D pixel positions. """ def __init__(self, in_h, in_w, out_features, max_freq=7): super().__init__() self.in_h = in_h self.in_w = in_w self.out_features = out_features # One FourierInterferenceField per output neuron # (share the spatial grid and freq_pairs, only A and φ differ) self.fields = nn.ModuleList([ FourierInterferenceField(in_h, in_w, max_freq) for _ in range(out_features) ]) self.bias = nn.Parameter(torch.zeros(out_features)) def get_weight_matrix(self): """Reconstruct full (out_features, in_h*in_w) weight matrix.""" rows = [field() for field in self.fields] # list of (H, W) W = torch.stack(rows, dim=0) # (out, H, W) return W.reshape(self.out_features, -1) # (out, H*W) def forward(self, x): W = self.get_weight_matrix() # (out, H*W) return F.linear(x, W, self.bias) # ───────────────────────────────────────────────────────────────── # Efficient batched version (shares freq computation across outputs) # ───────────────────────────────────────────────────────────────── class FourierInterferenceLinearFast(nn.Module): """ Same as above but vectorised across all output neurons at once. amplitudes : (out_features, n_freqs) phases : (out_features, n_freqs) Weight matrix computed in one batched cosine evaluation. """ def __init__(self, in_h, in_w, out_features, max_freq=7): super().__init__() self.in_h, self.in_w = in_h, in_w self.out_features = out_features n_freqs = (2 * max_freq + 1) ** 2 self.amplitudes = nn.Parameter(torch.randn(out_features, n_freqs) * 0.02) self.phases = nn.Parameter(torch.rand(out_features, n_freqs) * 2 * torch.pi) self.bias = nn.Parameter(torch.zeros(out_features)) freqs_1d = torch.arange(-max_freq, max_freq + 1) KX, KY = torch.meshgrid(freqs_1d, freqs_1d, indexing='ij') self.register_buffer('freq_pairs', torch.stack([KX.flatten(), KY.flatten()], dim=1).float()) gy = torch.linspace(0, 1, in_h) gx = torch.linspace(0, 1, in_w) GY, GX = torch.meshgrid(gy, gx, indexing='ij') self.register_buffer('spatial_grid', torch.stack([GY.flatten(), GX.flatten()], dim=1)) def get_weight_matrix(self): # base_phase: (n_pixels, n_freqs) base_phase = 2 * torch.pi * (self.spatial_grid @ self.freq_pairs.T) # phases broadcast: (1, n_pixels, n_freqs) + (out, 1, n_freqs) full_phase = base_phase.unsqueeze(0) + self.phases.unsqueeze(1) # W: (out, n_pixels) W = (torch.cos(full_phase) * self.amplitudes.unsqueeze(1)).sum(dim=-1) return W # (out_features, in_h*in_w) def forward(self, x): W = self.get_weight_matrix() return F.linear(x, W, self.bias) def get_weight_images(self): """Return (out_features, in_h, in_w) for visualisation.""" return self.get_weight_matrix().detach().reshape(self.out_features, self.in_h, self.in_w) def dominant_frequencies(self, digit_idx): """Return the top-10 most influential (kx, ky) for a given digit.""" importance = self.amplitudes[digit_idx].abs().detach() top_k = importance.topk(10) freqs = self.freq_pairs[top_k.indices] return [(int(f[0]), int(f[1]), v.item()) for f, v in zip(freqs, top_k.values)] # ───────────────────────────────────────────────────────────────── # Data # ───────────────────────────────────────────────────────────────── transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_loader = torch.utils.data.DataLoader( torchvision.datasets.MNIST('../data', train=True, download=True, transform=transform), batch_size=256, shuffle=True) test_loader = torch.utils.data.DataLoader( torchvision.datasets.MNIST('../data', train=False, download=True, transform=transform), batch_size=256, shuffle=False) # ───────────────────────────────────────────────────────────────── # Model # ───────────────────────────────────────────────────────────────── # Single Fourier interference layer (same depth as original SPDE experiment) model_fourier = nn.Sequential( nn.Flatten(), FourierInterferenceLinearFast(28, 28, out_features=10, max_freq=7) ) # For fair comparison: standard linear layer model_linear = nn.Sequential( nn.Flatten(), nn.Linear(784, 10) ) criterion = nn.CrossEntropyLoss() # ───────────────────────────────────────────────────────────────── # Training function # ───────────────────────────────────────────────────────────────── def train_model(model, name, epochs=10, lr=0.01): optimizer = torch.optim.Adam(model.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) history = {'loss': [], 'acc': []} for epoch in range(epochs): model.train() running_loss = 0.0 for images, labels in train_loader: optimizer.zero_grad() loss = criterion(model(images), labels) loss.backward() optimizer.step() running_loss += loss.item() * images.size(0) scheduler.step() model.eval() correct = total = 0 with torch.no_grad(): for images, labels in test_loader: _, pred = torch.max(model(images), 1) correct += (pred == labels).sum().item() total += labels.size(0) acc = 100 * correct / total loss = running_loss / len(train_loader.dataset) history['loss'].append(loss) history['acc'].append(acc) print(f"[{name}] Epoch {epoch+1:2d}/{epochs} loss={loss:.4f} acc={acc:.2f}%") return history # ───────────────────────────────────────────────────────────────── # Run training # ───────────────────────────────────────────────────────────────── print("=" * 55) print(" Fourier Interference Layer — MNIST") print("=" * 55) print(f"\nFourier layer params: {sum(p.numel() for p in model_fourier.parameters()):,}") print(f"Linear layer params: {sum(p.numel() for p in model_linear.parameters()):,}\n") history_fourier = train_model(model_fourier, "Fourier", epochs=10) print() history_linear = train_model(model_linear, "Linear", epochs=10) # ───────────────────────────────────────────────────────────────── # Visualisation # ───────────────────────────────────────────────────────────────── fourier_layer = model_fourier[1] linear_layer = model_linear[1] with torch.no_grad(): fourier_weights = fourier_layer.get_weight_images() # (10, 28, 28) linear_weights = linear_layer.weight.reshape(10, 28, 28) fig, axes = plt.subplots(3, 10, figsize=(18, 6)) for i in range(10): # Row 0: Fourier interference weights w = fourier_weights[i] axes[0, i].imshow(w, cmap='RdBu', vmin=-w.abs().max(), vmax=w.abs().max()) axes[0, i].set_title(f"{i}", fontsize=9) axes[0, i].axis('off') # Row 1: Standard linear weights w = linear_weights[i] axes[1, i].imshow(w, cmap='RdBu', vmin=-w.abs().max(), vmax=w.abs().max()) axes[1, i].axis('off') # Row 2: Amplitude spectrum — which frequencies matter for each digit? amps = fourier_layer.amplitudes[i].abs().detach().numpy() max_freq = fourier_layer.freq_pairs[:, 0].abs().max().int().item() n = 2 * max_freq + 1 axes[2, i].imshow(amps.reshape(n, n), cmap='hot', origin='lower') axes[2, i].axis('off') axes[0, 0].set_ylabel("Fourier\nweights", fontsize=8) axes[1, 0].set_ylabel("Linear\nweights", fontsize=8) axes[2, 0].set_ylabel("Amplitude\nspectrum", fontsize=8) acc_f = history_fourier['acc'][-1] acc_l = history_linear['acc'][-1] fig.suptitle( f"Fourier Interference vs Standard Linear | " f"Fourier acc: {acc_f:.1f}% Linear acc: {acc_l:.1f}%", fontsize=13, fontweight='bold' ) plt.tight_layout() plt.savefig('fourier_interference_weights.png', dpi=150, bbox_inches='tight') plt.show() # ── Training curves ── fig2, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) epochs_range = range(1, 11) ax1.plot(epochs_range, history_fourier['loss'], 'b-o', label='Fourier Interference') ax1.plot(epochs_range, history_linear['loss'], 'r-s', label='Standard Linear') ax1.set_xlabel('Epoch'); ax1.set_ylabel('Loss') ax1.set_title('Training Loss'); ax1.legend(); ax1.grid(alpha=0.3) ax2.plot(epochs_range, history_fourier['acc'], 'b-o', label='Fourier Interference') ax2.plot(epochs_range, history_linear['acc'], 'r-s', label='Standard Linear') ax2.set_xlabel('Epoch'); ax2.set_ylabel('Test Accuracy (%)') ax2.set_title('Test Accuracy'); ax2.legend(); ax2.grid(alpha=0.3) plt.suptitle('Fourier Interference Layer vs Standard Linear — Training Curves', fontsize=12, fontweight='bold') plt.tight_layout() plt.savefig('fourier_interference_curves.png', dpi=150, bbox_inches='tight') plt.show() # ── Dominant frequencies per digit ── print("\n── Dominant frequencies per digit (kx, ky, amplitude) ──") for d in range(10): top = fourier_layer.dominant_frequencies(d) freq_str = " ".join(f"({kx:+d},{ky:+d})={v:.3f}" for kx, ky, v in top[:5]) print(f" Digit {d}: {freq_str}") print(f"\nFinal test accuracy — Fourier: {acc_f:.2f}% | Linear: {acc_l:.2f}%")