import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms import numpy as np import math from tqdm import tqdm # ---------------------------------------------------------------------- # 1. Define a simple MLP for MNIST # ---------------------------------------------------------------------- class SimpleMLP(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(28*28, 128) self.fc2 = nn.Linear(128, 64) self.fc3 = nn.Linear(64, 10) def forward(self, x): x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x # ---------------------------------------------------------------------- # 2. Singularity generator (asymptotic zeros of the Riemann zeta function) # gamma_n = 2π * n / log(n+1) (shifted to match first zero ~14.13) # ---------------------------------------------------------------------- def generate_zeros(N): """Return a list of N 'zeros' gamma_n (imaginary parts) for n=1..N.""" zeros = [] for n in range(1, N+1): # Asymptotic formula with a small offset to avoid log(1)=0 gamma = 2 * math.pi * n / math.log(n + 1) zeros.append(gamma) return zeros # Pre‑compute once NUM_ZEROS = 500 ZEROS = generate_zeros(NUM_ZEROS) # sorted increasing ZERO_MIN, ZERO_MAX = ZEROS[0], ZEROS[-1] def find_nearest_zero(t, zeros): """Return (index, residual) for the zero closest to t.""" idx = np.searchsorted(zeros, t) if idx == 0: return 0, t - zeros[0] if idx == len(zeros): return len(zeros)-1, t - zeros[-1] left = zeros[idx-1] right = zeros[idx] if t - left < right - t: return idx-1, t - left else: return idx, t - right # ---------------------------------------------------------------------- # 3. Compression / decompression functions # ---------------------------------------------------------------------- def compress_tensor(tensor, zeros, zero_min, zero_max): """ Compress a PyTorch tensor using singularity encoding. Returns: list of (index, residual), plus original min/max for scaling. """ flat = tensor.detach().cpu().numpy().flatten() if flat.size == 0: return [], 0.0, 0.0 vmin, vmax = flat.min(), flat.max() # Avoid division by zero if all values equal if vmax == vmin: vmax = vmin + 1e-6 # Scale each value to [zero_min, zero_max] scaled = zero_min + (zero_max - zero_min) * (flat - vmin) / (vmax - vmin) compressed = [] for t in scaled: idx, resid = find_nearest_zero(t, zeros) compressed.append((idx, resid)) return compressed, vmin, vmax def decompress_tensor(compressed, vmin, vmax, zeros, zero_min, zero_max, original_shape): """Reconstruct a tensor from compressed representation.""" if not compressed: return torch.zeros(original_shape) flat_recon = [] for idx, resid in compressed: gamma = zeros[idx] t = gamma + resid # Inverse scaling val = vmin + (vmax - vmin) * (t - zero_min) / (zero_max - zero_min) flat_recon.append(val) recon_np = np.array(flat_recon, dtype=np.float32).reshape(original_shape) return torch.from_numpy(recon_np) def compress_model(model, zeros, zero_min, zero_max): """Compress all parameters of a model.""" compressed_data = {} for name, param in model.named_parameters(): comp, vmin, vmax = compress_tensor(param.data, zeros, zero_min, zero_max) compressed_data[name] = (comp, vmin, vmax, param.data.shape) return compressed_data def decompress_model(model, compressed_data, zeros, zero_min, zero_max): """Load decompressed parameters into the model.""" for name, (comp, vmin, vmax, shape) in compressed_data.items(): new_param = decompress_tensor(comp, vmin, vmax, zeros, zero_min, zero_max, shape) model.get_parameter(name).data.copy_(new_param) # ---------------------------------------------------------------------- # 4. Training and evaluation utilities # ---------------------------------------------------------------------- def train_one_epoch(model, device, train_loader, optimizer, epoch): model.train() correct = 0 total = 0 for batch_idx, (data, target) in enumerate(tqdm(train_loader, desc=f"Train Epoch {epoch}")): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.cross_entropy(output, target) loss.backward() optimizer.step() pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() total += target.size(0) acc = 100. * correct / total return acc def evaluate(model, device, test_loader): model.eval() correct = 0 total = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() total += target.size(0) return 100. * correct / total # ---------------------------------------------------------------------- # 5. Main test: train baseline -> compress -> decompress -> fine-tune # ---------------------------------------------------------------------- def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # MNIST data loaders transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) train_dataset = datasets.MNIST('../data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST('../data', train=False, transform=transform) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False) # Baseline model model = SimpleMLP().to(device) optimizer = optim.Adam(model.parameters(), lr=0.001) print("\n🔹 Training baseline model (5 epochs)...") for epoch in range(1, 6): train_acc = train_one_epoch(model, device, train_loader, optimizer, epoch) test_acc = evaluate(model, device, test_loader) print(f"Epoch {epoch}: Train Acc = {train_acc:.2f}%, Test Acc = {test_acc:.2f}%") baseline_acc = evaluate(model, device, test_loader) print(f"\n✅ Baseline test accuracy: {baseline_acc:.2f}%\n") # Compress the model print("🔸 Compressing model parameters using singularity encoding...") compressed = compress_model(model, ZEROS, ZERO_MIN, ZERO_MAX) # Estimate compression ratio (naive) total_params = sum(p.numel() for p in model.parameters()) bits_per_index = math.ceil(math.log2(NUM_ZEROS)) # ~9 bits for 500 zeros bits_per_residual = 16 # we can store residual as half float compressed_bits = total_params * (bits_per_index + bits_per_residual) original_bits = total_params * 32 # float32 ratio = compressed_bits / original_bits print(f"Compression ratio: {ratio:.2f} (store {bits_per_index}+{bits_per_residual} bits per weight)") # Decompress into a new model decompressed_model = SimpleMLP().to(device) decompress_model(decompressed_model, compressed, ZEROS, ZERO_MIN, ZERO_MAX) # Evaluate after decompression after_decomp_acc = evaluate(decompressed_model, device, test_loader) print(f"📉 Test accuracy after decompression (no fine‑tune): {after_decomp_acc:.2f}%") # Fine‑tune the decompressed model print("\n🔹 Fine‑tuning decompressed model (3 epochs)...") optimizer_ft = optim.Adam(decompressed_model.parameters(), lr=0.0005) for epoch in range(1, 4): train_acc = train_one_epoch(decompressed_model, device, train_loader, optimizer_ft, epoch) test_acc = evaluate(decompressed_model, device, test_loader) print(f"Fine-tune Epoch {epoch}: Train Acc = {train_acc:.2f}%, Test Acc = {test_acc:.2f}%") final_acc = evaluate(decompressed_model, device, test_loader) print(f"\n🎯 Final test accuracy after fine‑tuning: {final_acc:.2f}%") print(f" Baseline was: {baseline_acc:.2f}%") print(f" Accuracy recovered: {final_acc - baseline_acc:+.2f}%") if __name__ == "__main__": main()