import torch import torch.nn as nn import torch.nn.functional as F from torchvision import datasets, transforms import time import os # --------------------------------------------------------------------------- # 1. SUPERBOOLEAN SUPERPOSITION BASIS (Zero-Parameter Feature Extractor) # --------------------------------------------------------------------------- class SuperBooleanBasis(nn.Module): """ Probability-ASM / SuperBoolean feature expansion. Input: 784 MNIST pixels in [0,1] Output: 392 * 16 = 6272 features (all 16 two-input Boolean gates per pair). """ def __init__(self, in_features=784): super().__init__() assert in_features % 2 == 0 self.num_pairs = in_features // 2 def forward(self, x): a = x[:, 0::2] # (B, 392) b = x[:, 1::2] # (B, 392) gates = torch.stack([ torch.zeros_like(a), # 0 FALSE (1.0 - a) * (1.0 - b), # 1 NOR a * (1.0 - b), # 2 A AND NOT B 1.0 - a, # 3 NOT A (1.0 - a) * b, # 4 NOT A AND B 1.0 - b, # 5 NOT B a + b - 2.0 * a * b, # 6 XOR 1.0 - a * b, # 7 NAND a * b, # 8 AND 1.0 - a - b + 2.0 * a * b, # 9 XNOR b, # 10 B 1.0 - b + a * b, # 11 A OR NOT B a, # 12 A 1.0 - a + a * b, # 13 NOT A OR B a + b - a * b, # 14 OR torch.ones_like(a), # 15 TRUE ], dim=-1) # (B, 392, 16) return gates.view(x.size(0), -1) # (B, 6272) # --------------------------------------------------------------------------- # 2. CCT WEIGHT SUPERPOSITION (Ternary Straight-Through Estimator) # --------------------------------------------------------------------------- class SuperPositionLinear(nn.Module): """ Ternary linear layer with a Straight-Through Estimator. FIX (was the main accuracy killer): - Old code initialized latent weights with std=0.01 but collapsed with fixed thresholds at +/-0.5, so EVERY weight collapsed to 0. The layer output was just its (zero) bias, and because the forward pass uses the collapsed weights, no gradient could flow backward through the layer. The latent field could never escape the dead zone. - New code uses the Ternary Weight Networks formulation: delta = 0.75 * mean(|w_latent|) (scale-relative threshold) alpha = mean(|w_latent|) over the non-zero weights (scaling factor) so a healthy fraction of weights is always +/-alpha, and activations keep a sane magnitude. The discrete state set is still {-1, 0, +1} (times one scalar alpha per layer). """ def __init__(self, in_features, out_features, bias=True, ternary=True): super().__init__() self.ternary = ternary # FIX: sensible init scale (~1/sqrt(fan_in)) instead of 0.01 self.weight_latent = nn.Parameter( torch.randn(out_features, in_features) / (in_features ** 0.5)) if bias: self.bias = nn.Parameter(torch.zeros(out_features)) else: self.register_parameter('bias', None) def collapse(self): """CCT: collapse the superposition to discrete states for measurement.""" if self.ternary: w = self.weight_latent delta = 0.75 * w.abs().mean() ternary = (w > delta).float() - (w < -delta).float() n_nonzero = ternary.abs().sum().clamp(min=1.0) alpha = (w.abs() * ternary.abs()).sum() / n_nonzero return alpha * ternary else: return torch.sign(self.weight_latent) def forward(self, x): w_collapsed = self.collapse() # Straight-Through Estimator: forward uses discrete weights, # gradient flows to the continuous latent field. w_ste = self.weight_latent + (w_collapsed - self.weight_latent).detach() return F.linear(x, w_ste, self.bias) def effective_memory_bytes(self): bits = self.weight_latent.numel() * 1.585 if self.bias is not None: bits += self.bias.numel() * 16.0 return int(bits / 8) # --------------------------------------------------------------------------- # 3. CCT CLASSIFIER WITH ENTROPY-AWARE EARLY COLLAPSE # --------------------------------------------------------------------------- class CCT_MNISTClassifier(nn.Module): def __init__(self, entropy_threshold=0.40): super().__init__() self.entropy_threshold = entropy_threshold self.superboolean = SuperBooleanBasis(784) sb_dim = 392 * 16 # 6272 # CHEAP measurement head self.cheap_head = nn.Linear(sb_dim, 10) # DEEP measurement head self.deep = nn.Sequential( SuperPositionLinear(sb_dim, 100, ternary=True), nn.ReLU(), SuperPositionLinear(100, 100, ternary=True), nn.ReLU(), nn.Linear(100, 10) ) self.register_buffer('cheap_collapses', torch.tensor(0)) self.register_buffer('total_samples', torch.tensor(0)) def forward(self, x, use_cct=True): x = x.view(x.size(0), -1) x = torch.clamp(x, 0.0, 1.0) sb_features = self.superboolean(x) cheap_logits = self.cheap_head(sb_features) # FIX: during training, return BOTH heads so both get a loss signal. # Previously cheap_head was never trained, so the entropy check at # test time was computed on a randomly initialized head, and early # collapses returned random predictions. if self.training: deep_logits = self.deep(sb_features) return cheap_logits, deep_logits if use_cct: probs = F.softmax(cheap_logits, dim=1) entropy = -(probs * torch.log(probs + 1e-12)).sum(dim=1) # Batch-wise conditional collapse (SIMD-friendly). if entropy.mean() < self.entropy_threshold: self.cheap_collapses += x.size(0) self.total_samples += x.size(0) return cheap_logits self.total_samples += x.size(0) return self.deep(sb_features) def cct_stats(self): if self.total_samples.item() == 0: return 0.0 return 100.0 * self.cheap_collapses.item() / self.total_samples.item() def print_memory_footprint(self): total_latent = sum(p.numel() * p.element_size() for p in self.parameters()) total_buffer = sum(b.numel() * b.element_size() for b in self.buffers()) effective = 0 for m in self.modules(): if isinstance(m, SuperPositionLinear): effective += m.effective_memory_bytes() effective += sum(p.numel() * 4 for n, p in self.named_parameters() if 'cheap_head' in n or 'deep.4' in n) print(f"Latent RAM footprint: {(total_latent + total_buffer) / 1024:.2f} KB") print(f"Effective info payload: {effective / 1024:.2f} KB") print(f"(Ternary layers use ~1.58 bits/weight vs 32 bits latent)") # --------------------------------------------------------------------------- # 4. TRAINING & INFERENCE LOOPS # --------------------------------------------------------------------------- def train_epoch(model, device, train_loader, optimizer, epoch): model.train() total_loss = 0.0 correct = 0 t0 = time.time() for data, target in train_loader: data, target = data.to(device), target.to(device) for i in range(10): optimizer.zero_grad() # FIX: removed the hardcoded CPU autocast (it was applied even on # CUDA, and bf16 on CPU only hurts a model this small). cheap_logits, deep_logits = model(data) # Joint loss: deep head is primary, cheap head gets its own signal # so CCT entropy routing is meaningful at inference time. loss = F.cross_entropy(deep_logits, target) \ + 0.5 * F.cross_entropy(cheap_logits, target) loss.backward() optimizer.step() if i==0: total_loss += loss.item() * data.size(0) pred = deep_logits.argmax(dim=1) correct += pred.eq(target).sum().item() elapsed = time.time() - t0 avg_loss = total_loss / len(train_loader.dataset) acc = 100.0 * correct / len(train_loader.dataset) print(f"Epoch {epoch}: Loss={avg_loss:.4f} Acc={acc:.2f}% Time={elapsed:.2f}s") return elapsed def test(model, device, test_loader, use_cct=True): model.eval() loss = 0.0 correct = 0 t0 = time.time() with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data, use_cct=use_cct) loss += F.cross_entropy(output, target, reduction='sum').item() pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() elapsed = time.time() - t0 avg_loss = loss / len(test_loader.dataset) acc = 100.0 * correct / len(test_loader.dataset) return avg_loss, acc, elapsed # --------------------------------------------------------------------------- # 5. MAIN # --------------------------------------------------------------------------- def main(): torch.backends.cudnn.benchmark = True if not torch.cuda.is_available(): torch.set_num_threads(min(4, os.cpu_count() or 1)) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Device: {device}") transform = transforms.Compose([transforms.ToTensor()]) train_ds = datasets.MNIST('data', train=True, download=True, transform=transform) test_ds = datasets.MNIST('data', train=False, transform=transform) BATCH_SIZE = 512 train_loader = torch.utils.data.DataLoader( train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2) test_loader = torch.utils.data.DataLoader( test_ds, batch_size=BATCH_SIZE * 2, shuffle=False, num_workers=2) model = CCT_MNISTClassifier(entropy_threshold=0.35).to(device) print("\n--- Model Architecture ---") print(model) print("\n--- Memory Analysis ---") model.print_memory_footprint() if hasattr(torch, 'compile'): print("\nApplying torch.compile() ...") model = torch.compile(model) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) print("\n--- Training ---") total_train_time = 0.0 for epoch in range(1, 100): total_train_time += train_epoch(model, device, train_loader, optimizer, epoch) print("\n--- Testing with CCT Early Collapse ---") test_loss, test_acc, test_time = test(model, device, test_loader, use_cct=True) print(f"\nTest Loss: {test_loss:.4f}") print(f"Test Accuracy: {test_acc:.2f}%") print(f"Inference Time: {test_time:.3f}s for {len(test_loader.dataset)} samples") print(f"CCT Cheap Collapse Rate: {model.cct_stats():.1f}% of samples collapsed cheaply") print(f"Total Training Time: {total_train_time:.2f}s") # Sanity check: accuracy of the deep path alone (CCT off) _, deep_acc, _ = test(model, device, test_loader, use_cct=False) print(f"Deep-path-only Accuracy (CCT off): {deep_acc:.2f}%") if __name__ == '__main__': main()