import numpy as np from scipy.special import softmax class CCTPeriodicLayer: """ Efficient periodic layer: y = sum_i A_i * sin(2π f_i x + φ_i) All operations vectorised over batch and features. """ def __init__(self, in_features, n_functions): self.in_features = in_features self.n_functions = n_functions # Trainable parameters self.amplitudes = np.random.randn(n_functions) * 0.01 self.frequencies = np.random.rand(n_functions) * 2.0 self.phases = np.random.rand(n_functions) * 2 * np.pi # Precompute frequency matrix for all input dimensions # Each function has its own frequency per input dimension (learnable) # But to keep parameter count low, we use the same freq for all dims # and then a small linear projection? Instead, we'll keep per-dim freq as a single scalar. # Actually: sin(2π f * x + φ) where x is a scalar per input dimension. # For each function, we sum over input dimensions: sum_j sin(2π f * x_j + φ_j) # That would require n_functions * in_features phases -> too many parameters. # Better: standard Fourier feature: gamma(x) = [sin(2π f_k x), cos(2π f_k x)] concatenated. # But we want output of shape (batch, n_functions), not (batch, 2*n_functions*in_features). # Let's use the "random Fourier features" approach: # For each function, generate a random weight vector w (length in_features) # Then compute sin(w·x + phase) . That gives one output per function. # Parameters: w (n_functions, in_features) -> large. We want minimal. # Alternative: use 1D frequencies per function, but average over input dimensions # after applying per-dim phase shifts? That becomes messy. # SOLUTION: Use a small trainable linear projection from in_features to n_functions, # then apply periodic activation. This adds n_functions * in_features parameters, # which defeats minimalism. So we stick with the original "sum over dims" but # with vectorised computation and BACKPROP. # We'll implement the original idea (sum over input dims) but fully vectorised # and with analytic gradients. Parameter count = 3 * n_functions. def forward(self, x): """ x: (batch, in_features) returns (batch, n_functions) """ batch, d = x.shape # Expand: (batch, 1, d) * (1, n_f, d) not needed, we sum over d. # Compute angle for each (batch, function, dim) # angle = 2π * f * x + phase # Then sum over dims: output[b,f] = sum_d A_f * sin(angle[b,f,d]) # Precompute frequency matrix: for each function and each input dim, we need freq. # We'll use the same freq per function for all dims, plus a small per-dim modulation (optional). # To keep parameter count low, freq is per function (n_f) only. # angle = 2π * freq_f * x_bd + phase_f + pi * x_bd (as in original, but phase_f only). # Let's simplify: angle = 2π * freq_f * x_bd + phase_f # Then sum over d. freq = self.frequencies[:, None] # (n_f, 1) x_exp = x[None, :, :] # (1, batch, d) angle = 2 * np.pi * freq * x_exp + self.phases[:, None, None] # (n_f, batch, d) sin_vals = np.sin(angle) # (n_f, batch, d) # Sum over input dimension out = np.sum(self.amplitudes[:, None, None] * sin_vals, axis=2) # (n_f, batch) return out.T # (batch, n_f) def backward(self, x, grad_output): """ grad_output: (batch, n_f) gradient from upstream returns gradients w.r.t amplitudes, frequencies, phases """ batch, d = x.shape n_f = self.n_functions freq = self.frequencies[:, None] # (n_f,1) x_exp = x[None, :, :] # (1, batch, d) angle = 2 * np.pi * freq * x_exp + self.phases[:, None, None] # (n_f, batch, d) sin_vals = np.sin(angle) cos_vals = np.cos(angle) # Gradient w.r.t amplitudes: dL/dA = sum_{b,d} sin(angle) * grad_output[b,f] # grad_output shape (batch, n_f) -> expand to (n_f, batch, 1) g_out_exp = grad_output.T[:, :, None] # (n_f, batch, 1) dA = np.sum(g_out_exp * sin_vals, axis=(1,2)) # (n_f,) # Gradient w.r.t phases: dL/dphase = sum_{b,d} A * cos(angle) * grad_output[b,f] dPhase = np.sum(self.amplitudes[:, None, None] * g_out_exp * cos_vals, axis=(1,2)) # Gradient w.r.t frequencies: dL/df = sum_{b,d} A * cos(angle) * (2π * x) * grad_output[b,f] df = np.sum(self.amplitudes[:, None, None] * g_out_exp * cos_vals * (2 * np.pi * x_exp), axis=(1,2)) return dA, dPhase, df class FastCCTClassifier: """ Minimal periodic classifier with one periodic hidden layer + linear output. Trainable parameters: 3*n_functions (periodic) + (n_functions+1)*n_classes (linear output) Still much smaller than MLP but enough to learn. """ def __init__(self, input_size, n_functions, n_classes): self.input_size = input_size self.n_functions = n_functions self.n_classes = n_classes # Periodic layer self.periodic = CCTPeriodicLayer(input_size, n_functions) # Output linear layer (adds moderate parameters: n_functions * n_classes) self.W_out = np.random.randn(n_functions, n_classes) * 0.01 self.b_out = np.zeros(n_classes) # For entropy tracking (optional) self.H = 1.0 def forward(self, x): # Periodic mapping z = self.periodic.forward(x) # (batch, n_functions) # Linear output logits = z @ self.W_out + self.b_out # (batch, n_classes) # Softmax exp_logits = np.exp(logits - np.max(logits, axis=1, keepdims=True)) probs = exp_logits / np.sum(exp_logits, axis=1, keepdims=True) # Entropy self.H = -np.mean(np.sum(probs * np.log(probs + 1e-9), axis=1)) return probs, z def backward(self, x, y_onehot, z, probs): batch = x.shape[0] # Gradient w.r.t logits (cross-entropy) dlogits = probs - y_onehot # (batch, n_classes) # Gradients for output layer dW = z.T @ dlogits / batch db = np.mean(dlogits, axis=0) # Gradient w.r.t periodic layer output dz = dlogits @ self.W_out.T # (batch, n_functions) # Backprop through periodic layer dA, dPhase, df = self.periodic.backward(x, dz) return dA, dPhase, df, dW, db def update(self, x, y_onehot, lr=0.01): probs, z = self.forward(x) dA, dPhase, df, dW, db = self.backward(x, y_onehot, z, probs) # Update periodic parameters self.periodic.amplitudes -= lr * dA self.periodic.phases -= lr * dPhase self.periodic.frequencies -= lr * df # Keep amplitudes positive and phases bounded self.periodic.amplitudes = np.maximum(self.periodic.amplitudes, 0.001) self.periodic.phases = self.periodic.phases % (2 * np.pi) # Update output layer self.W_out -= lr * dW self.b_out -= lr * db return probs def predict(self, x): probs, _ = self.forward(x) return np.argmax(probs, axis=1) def score(self, x, y): acc = np.mean(self.predict(x) == y) return acc # ================================ # DEMO on synthetic classification # ================================ if __name__ == "__main__": print("Fast & Learnable CCT-ODE Classifier") print("====================================") # Create synthetic spiral data (nonlinear) from sklearn.datasets import make_classification X, y = make_classification(n_samples=2000, n_features=20, n_classes=3, n_informative=15, random_state=42) X = (X - X.mean(axis=0)) / (X.std(axis=0) + 1e-9) # normalize split = 1600 X_train, X_test = X[:split], X[split:] y_train, y_test = y[:split], y[split:] # Model with 128 periodic functions (384 parameters) + output layer 128*3+3 = 387 # Total = 771 parameters, still tiny compared to MLP (~20*200??) model = FastCCTClassifier(input_size=20, n_functions=128, n_classes=3) print(f"Trainable parameters: periodic = {3*128}, output = {128*3+3}, total = {3*128 + 128*3+3}") # Training loop batch_size = 64 epochs = 200 for epoch in range(epochs): idx = np.random.choice(len(X_train), batch_size) X_batch = X_train[idx] y_batch = y_train[idx] y_onehot = np.eye(3)[y_batch] model.update(X_batch, y_onehot, lr=0.01) if epoch % 20 == 0: train_acc = model.score(X_train[:200], y_train[:200]) test_acc = model.score(X_test, y_test) print(f"Epoch {epoch} | Train acc: {train_acc:.3f} | Test acc: {test_acc:.3f} | Entropy: {model.H:.3f}") print("\nFinal test accuracy:", model.score(X_test, y_test))