import numpy as np from scipy.io.wavfile import read # ============================================================ # Load MNIST data (wav format as provided) # ============================================================ X_train = read('../X_train.wav')[1].reshape(-1, 784) y_train = (read('../y_train.wav')[1] * 9).astype(int) X_test = read('../X_test.wav')[1].reshape(-1, 784) y_test = (read('../y_test.wav')[1] * 9).astype(int) X20 = X_test[:1000] yt20 = y_test[:1000] # ============================================================ # Im2Col helper: unfold image patches for conv-as-matmul # ============================================================ def im2col(X, kernel_size, stride=1, padding=1): """Unfold image into column matrix for convolution-as-matmul. Args: X: (B, H, W, C) image tensor kernel_size: size of conv kernel stride: conv stride padding: zero padding around image Returns: cols: (B, H_out*W_out, K*K*C) unfolded patches """ B, H, W, C = X.shape # Pad input X_padded = np.pad(X, ((0, 0), (padding,)*2, (padding,)*2, (0, 0))) out_h = (H + 2*padding - kernel_size) // stride + 1 out_w = (W + 2*padding - kernel_size) // stride + 1 cols = np.zeros((B, out_h * out_w, kernel_size * kernel_size * C)) for y in range(out_h): for x in range(out_w): patch = X_padded[:, y*stride : y*stride + kernel_size, x*stride : x*stride + kernel_size, :] cols[:, y*out_w + x, :] = patch.reshape(B, -1) return cols def col2im(cols, input_shape, kernel_size, stride=1, padding=1): """Reverse im2col: fold columns back into image.""" B, H, W, C = input_shape out_h = (H + 2*padding - kernel_size) // stride + 1 out_w = (W + 2*padding - kernel_size) // stride + 1 X_padded = np.zeros((B, H + 2*padding, W + 2*padding, C)) idx = 0 for y in range(out_h): for x in range(out_w): patch = cols[:, idx, :].reshape(B, kernel_size, kernel_size, C) X_padded[:, y*stride:y*stride+kernel_size, x*stride:x*stride+kernel_size, :] += patch idx += 1 if padding > 0: return X_padded[:, padding:-padding, padding:-padding, :] return X_padded # ============================================================ # ConvApproxMLP: MLP with Big/Little W1 weight matrices # ============================================================ class ConvApproxMLP: """MLP that approximates Conv2D using Im2Col transformation. - W1_big: Large weight matrix (H*W*C, hidden) - standard MLP weights - W1_little: Small weight matrix (K*K*C, hidden) - conv kernel approximation Training learns both, blending captures both local patterns (Conv) and global patterns (MLP). """ def __init__(self, input_size=784, hidden_size=100, output_size=10, kernel_size=3, learning_rate=None, conv_alpha=0.5): """ Args: input_size: 784 for MNIST (28x28) hidden_size: number of hidden units output_size: 10 for MNIST digits kernel_size: size of local receptive field for little W1 learning_rate: array of 3 rates [lr_big, lr_little, lr_out] conv_alpha: blend weight for big-path vs little-path (1=all big, 0=all little) """ self.input_size = input_size self.hidden_size = hidden_size self.output_size = output_size self.kernel_size = kernel_size self.conv_alpha = conv_alpha # Set learning rates if learning_rate is None: learning_rate = np.array([0.01, 0.005, 0.01], dtype=np.float32) self.learning_rate = np.asarray(learning_rate, dtype=np.float32) if self.learning_rate.shape != (3,): raise ValueError("learning_rate must be a length-3 array: [lr_big, lr_little, lr_out]") # Image dimensions (assume square) self.img_size = int(np.sqrt(input_size)) # ---- BIG W1: Full dense MLP weights (each pixel → hidden) ---- # Shape: (784, hidden_size) self.W1_big = np.random.randn(input_size, hidden_size) * np.sqrt(2.0 / input_size) self.b1_big = np.zeros((1, hidden_size)) # ---- LITTLE W1: Small conv-like weights (local patches → hidden) ---- # Shape: (kernel_size * kernel_size, hidden_size) little_fan_in = kernel_size * kernel_size self.W1_little = np.random.randn(little_fan_in, hidden_size) * np.sqrt(2.0 / little_fan_in) self.b1_little = np.zeros((1, hidden_size)) # ---- OUTPUT LAYER ---- self.W2 = np.random.randn(hidden_size, output_size) * np.sqrt(2.0 / hidden_size) self.b2 = np.zeros((1, output_size)) # For im2col with padding (to keep spatial size) self.padding = kernel_size // 2 def _reshape_to_image(self, X): """Convert flat (B, 784) to image (B, 28, 28, 1)""" if len(X.shape) == 2: B = X.shape[0] return X.reshape(B, self.img_size, self.img_size, 1) return X def _forward_big(self, X): """Standard MLP forward: W1_big @ x + b1_big""" X_flat = X.reshape(-1, self.input_size) z1 = np.dot(X_flat, self.W1_big) + self.b1_big a1 = self.relu(z1) self.z1_big = z1 return a1 def _forward_little(self, X): """Conv approximation: im2col(x) @ W1_little + b1_little""" # Reshape to image X_img = self._reshape_to_image(X) B = X_img.shape[0] # Unfold into columns: (B, spatial_out, K*K*C) X_cols = im2col(X_img, self.kernel_size, stride=1, padding=self.padding) # Reshape for matmul: (B*K*K*C, hidden) -> but we do per-sample # Actually: each spatial position uses same W1_little (weight sharing!) # X_cols shape: (B, H_out*W_out, K*K*C) # Reshape W1_little: (K*K, hidden) -> tiled for broadcasting # Output: sum over local patch weighted by W1_little # This is equivalent to: a1_spatial[b, pos, h] = sum_k(X_patch[b, pos, k] * W1_little[k, h]) # Simple approach: im2col gives (B, spatial, K*K*C), W1_little is (K*K*C, hidden) # Result: (B, spatial, hidden) then pool/flatten spatial = X_cols.shape[1] # For each batch, compute conv-like operation # X_cols: (B, spatial, K*K) @ W1_little: (K*K, hidden) -> (B, spatial, hidden) z1_spatial = np.dot(X_cols, self.W1_little) # (B, spatial, hidden) # Add bias (broadcast to spatial dimension) z1_spatial = z1_spatial + self.b1_little # (B, spatial, hidden) # Apply activation a1_spatial = self.relu(z1_spatial) # (B, spatial, hidden) self.z1_spatial = z1_spatial # Pool spatially (mean) to get (B, hidden) a1 = np.mean(a1_spatial, axis=1) # (B, hidden) return a1 def relu(self, x): return np.maximum(0, x) def relu_derivative(self, x): return np.where(x > 0, 1, 0) def softmax(self, x): exp_x = np.exp(x - np.max(x, axis=1, keepdims=True)) return exp_x / np.sum(exp_x, axis=1, keepdims=True) def forward(self, X): """Combined forward: blend big (MLP) and little (Conv) representations""" a1_big = self._forward_big(X) a1_little = self._forward_little(X) # Blend based on conv_alpha # conv_alpha=1: pure big path, conv_alpha=0: pure little path self.a1 = self.conv_alpha * a1_big + (1 - self.conv_alpha) * a1_little # Store intermediate values for backprop self.a1_big = a1_big self.a1_little = a1_little # Store original input shapes self.X_img = self._reshape_to_image(X) self.X_flat = X.reshape(-1, self.input_size) # Output layer z2 = np.dot(self.a1, self.W2) + self.b2 output = self.softmax(z2) return output def compute_loss(self, y_true, y_pred): """Cross-entropy loss""" m = y_true.shape[0] loss = -np.sum(y_true * np.log(y_pred + 1e-9)) / m return loss def backward(self, X, y_true, y_pred): """Backprop through combined big/little architecture""" m = y_true.shape[0] # ---- Output layer gradients ---- dz2 = y_pred - y_true # (m, 10) dW2 = np.dot(self.a1.T, dz2) / m db2 = np.sum(dz2, axis=0, keepdims=True) / m # ---- Hidden layer gradients (shared by big and little paths) ---- da1 = np.dot(dz2, self.W2.T) # (m, hidden) # ---- Big path gradients ---- # Gradients through MLP path da1_big = self.conv_alpha * da1 # dz1_big = da1_big * relu'(z1_big) - but we only stored a1_big # Approximate: use a1_big > 0 as proxy for activation dz1_big_approx = da1_big * self.relu_derivative(self.z1_big) dW1_big = np.dot(self.X_flat.T, dz1_big_approx) / m db1_big = np.sum(dz1_big_approx, axis=0, keepdims=True) / m # ---- Little path gradients (Conv path) ---- da1_little = (1 - self.conv_alpha) * da1 # For the little path, gradients need to flow back through spatial pooling # Each spatial position contributed equally to a1_little (mean pooling) da1_spatial = da1_little / (self.X_img.shape[1] * self.X_img.shape[2]) # (B, hidden) # Broadcast to all spatial positions da1_spatial_full = np.expand_dims(da1_spatial, 1) # (B, 1, hidden) da1_spatial_full = np.tile(da1_spatial_full, [1, self.X_img.shape[1] * self.X_img.shape[2], 1]) # dz1_spatial = da1_spatial * relu'(z1_spatial) # We didn't store z1_spatial, approximate using a1_spatial dz1_spatial = da1_spatial_full * self.relu_derivative(self.z1_spatial) # Gradient w.r.t. W1_little: im2col(X)^T @ dz1_spatial X_cols = im2col(self.X_img, self.kernel_size, stride=1, padding=self.padding) # Reshape dz1_spatial for matmul dz1_spatial_2d = dz1_spatial.reshape(-1, self.hidden_size) # (B*spatial, hidden) X_cols_2d = X_cols.reshape(-1, self.kernel_size * self.kernel_size) # (B*spatial, K*K) dW1_little = np.dot(X_cols_2d.T, dz1_spatial_2d) / m db1_little = np.sum(dz1_spatial, axis=(0, 1), keepdims=True) / m return dW1_big, db1_big, dW1_little, db1_little, dW2, db2 def _forward_little_raw(self, X): """Raw forward for little path (no caching)""" X_img = self._reshape_to_image(X) X_cols = im2col(X_img, self.kernel_size, stride=1, padding=self.padding) z1_spatial = np.dot(X_cols, self.W1_little) + self.b1_little a1_spatial = self.relu(z1_spatial) return a1_spatial def update(self, X, y_true): """Single training step""" y_pred = self.forward(X) grads = self.backward(X, y_true, y_pred) dW1_big, db1_big, dW1_little, db1_little, dW2, db2 = grads # Update big W1 self.W1_big -= self.learning_rate[0] * dW1_big self.b1_big -= self.learning_rate[0] * db1_big # Update little W1 self.W1_little -= self.learning_rate[1] * dW1_little self.b1_little -= self.learning_rate[1] * db1_little.reshape(1, -1) # Update output layer self.W2 -= self.learning_rate[2] * dW2 self.b2 -= self.learning_rate[2] * db2 def predict(self, X): probabilities = self.forward(X) return np.argmax(probabilities, axis=1) def score(self, X, y_true): y = self.predict(X) return np.mean(y == y_true) def summary(self): """Print model architecture summary""" big_params = self.W1_big.size + self.b1_big.size little_params = self.W1_little.size + self.b1_little.size out_params = self.W2.size + self.b2.size print(f"ConvApproxMLP Summary:") print(f" Input: {self.input_size}") print(f" Hidden: {self.hidden_size}") print(f" Output: {self.output_size}") print(f" Kernel size: {self.kernel_size}") print(f" conv_alpha: {self.conv_alpha}") print(f" W1_big (MLP): {self.W1_big.shape} = {big_params:,} params") print(f" W1_little (Conv): {self.W1_little.shape} = {little_params:,} params") print(f" W2 (output): {self.W2.shape} = {out_params:,} params") print(f" Total params: {big_params + little_params + out_params:,}") print(f" Compression vs MLP: {big_params / (big_params + little_params):.1%} of original") # ============================================================ # Initialize and train # ============================================================ # Learning rates for [big path, little path, output layer] # The original rates were so small that the model stayed near chance. learning_rates = np.array([0.05, 0.01, 0.05], dtype=np.float32) # Initialize ConvApproxMLP f = ConvApproxMLP( input_size=784, hidden_size=100, output_size=10, kernel_size=3, learning_rate=learning_rates, conv_alpha=0.9 # Keep the little path, but let the learnable big path dominate early training ) f.summary() # ============================================================ # Training loop (same structure as original) # ============================================================ print("\nTraining:") i = 0 while True: # Random minibatch idx = np.random.randint(0, 60000, 100) X = X_train[idx] yt = y_train[idx] # Onehot encode labels y_onehot = np.eye(10)[yt] # Training step f.update(X, y_onehot) # Print progress if i % 50 == 0: acc = f.score(X, yt) y_pred = f.forward(X) loss = f.compute_loss(y_onehot, y_pred) print(f" Iter {i:5d} | Loss: {loss:.4f} | Train Acc: {acc:.4f}") # Evaluate on test set periodically if i > 0 and i % 500 == 0: test_acc = f.score(X_test, y_test) print(f" Iter {i:5d} | Test Acc: {test_acc:.4f}") i += 1 # Stop condition if i >= 5000: break # ============================================================ # Final evaluation # ============================================================ print("\n" + "="*50) print("FINAL EVALUATION") print("="*50) final_train_acc = f.score(X_train[:5000], y_train[:5000]) final_test_acc = f.score(X_test, y_test) print(f"Training accuracy (first 5000): {final_train_acc:.4f}") print(f"Test accuracy: {final_test_acc:.4f}") # Show weight matrix statistics print(f"\nWeight matrices:") print(f" W1_big range: [{f.W1_big.min():.4f}, {f.W1_big.max():.4f}]") print(f" W1_little range: [{f.W1_little.min():.4f}, {f.W1_little.max():.4f}]") print(f" W2 range: [{f.W2.min():.4f}, {f.W2.max():.4f}]")