File size: 4,492 Bytes
944a442 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | #!/usr/bin/env python3
"""Construct neurons to match two Jacobian boundaries at x1 ≈ -0.665558 and x2 ≈ 0.594541."""
import torch
import math
import os
from safetensors.torch import save_file
# Target boundaries and slopes
boundary_x1 = -0.665558
boundary_x2 = 0.594541
left_slope = 11.0
mid_slope = 1.0
right_slope = 0.5
# Right equation: y = 0.5 * x - 0.2489895
right_eq = lambda x: 0.5 * x - 0.2489895
# Middle equation (continuous at boundary_x2):
c_mid = right_eq(boundary_x2) - mid_slope * boundary_x2
mid_eq = lambda x: mid_slope * x + c_mid
# Left equation (continuous at boundary_x1):
c_left = mid_eq(boundary_x1) - left_slope * boundary_x1
left_eq = lambda x: left_slope * x + c_left
print("Constructing neurons to match:")
print(f" Boundary 1 at x = {boundary_x1}")
print(f" Boundary 2 at x = {boundary_x2}")
print(f" Left slope = {left_slope:4.1f}, range: x < {boundary_x1:.6f}")
print(f" Middle slope = {mid_slope:4.1f}, range: {boundary_x1:.6f} < x < {boundary_x2:.6f}")
print(f" Right slope = {right_slope:4.1f}, range: x > {boundary_x2:.6f}")
# Use native PyTorch tensors to avoid redundant casting overhead
W1 = torch.zeros((8, 1), dtype=torch.float32)
b1 = torch.zeros(8, dtype=torch.float32)
W2 = torch.zeros((1, 8), dtype=torch.float32)
b2 = torch.zeros(1, dtype=torch.float32)
# Neuron 0: Always active pure slope carrier
# FIX: Use a robustly large bias so the neuron doesn't turn off during extreme negative activation outliers
W1[0, 0] = 1.0
b1[0] = 100.0 # Guarantees the carrier stays active for x > -100.0
W2[0, 0] = right_slope
# Neuron 1: Active left of boundary_x1 (adds left_slope - mid_slope = 10.0 to slope)
W1[1, 0] = -1.0
b1[1] = boundary_x1
W2[0, 1] = -(left_slope - mid_slope)
# Neuron 2: Active left of boundary_x2 (adds mid_slope - right_slope = 0.5 to slope)
W1[2, 0] = -1.0
b1[2] = boundary_x2
W2[0, 2] = -(mid_slope - right_slope)
# Neurons 3-7: Inactive (zero weights)
# Calculate exact b2 so that the function matches target_y at boundary_x2
target_y = right_eq(boundary_x2)
neuron0_out = W2[0, 0] * (W1[0, 0] * boundary_x2 + b1[0])
b2[0] = target_y - neuron0_out
print("\nConstructed neuron weights:")
print(f"W1:\n{W1.numpy()}")
print(f"b1: {b1.numpy()}")
print(f"W2:\n{W2.numpy()}")
print(f"b2: {b2.numpy()}")
# Verify the construction natively
def mlp_forward(x):
x_t = torch.tensor([[x]], dtype=torch.float32)
h = x_t @ W1.T + b1
h = torch.relu(h)
y = h @ W2.T + b2
return y.item()
print("\n" + "=" * 60)
print("Automated Verification:")
print("=" * 60)
# Test at Boundary 1
y_b1 = mlp_forward(boundary_x1)
expected_y_b1 = mid_eq(boundary_x1)
assert math.isclose(y_b1, expected_y_b1, rel_tol=1e-4, abs_tol=1e-5), f"Boundary 1 mismatch! Expected {expected_y_b1}, got {y_b1}"
print(f"✓ Boundary 1 (x = {boundary_x1}) matched")
# Test at Boundary 2
y_b2 = mlp_forward(boundary_x2)
expected_y_b2 = right_eq(boundary_x2)
assert math.isclose(y_b2, expected_y_b2, rel_tol=1e-4, abs_tol=1e-5), f"Boundary 2 mismatch! Expected {expected_y_b2}, got {y_b2}"
print(f"✓ Boundary 2 (x = {boundary_x2}) matched")
# Test left slope
x_left1, x_left2 = boundary_x1 - 0.10, boundary_x1 - 0.05
slope_left = (mlp_forward(x_left2) - mlp_forward(x_left1)) / (x_left2 - x_left1)
assert math.isclose(slope_left, left_slope, rel_tol=1e-4, abs_tol=1e-5), f"Left slope mismatch! Expected {left_slope}, got {slope_left}"
print(f"✓ Left slope matched: {slope_left:.4f}")
# Test middle slope
x_mid1, x_mid2 = (boundary_x1 + boundary_x2) / 2, ((boundary_x1 + boundary_x2) / 2) + 0.05
slope_mid = (mlp_forward(x_mid2) - mlp_forward(x_mid1)) / (x_mid2 - x_mid1)
assert math.isclose(slope_mid, mid_slope, rel_tol=1e-4, abs_tol=1e-5), f"Middle slope mismatch! Expected {mid_slope}, got {slope_mid}"
print(f"✓ Middle slope matched: {slope_mid:.4f}")
# Test right slope
x_right1, x_right2 = boundary_x2 + 0.05, boundary_x2 + 0.10
slope_right = (mlp_forward(x_right2) - mlp_forward(x_right1)) / (x_right2 - x_right1)
assert math.isclose(slope_right, right_slope, rel_tol=1e-4, abs_tol=1e-5), f"Right slope mismatch! Expected {right_slope}, got {slope_right}"
print(f"✓ Right slope matched: {slope_right:.4f}")
os.makedirs("test_mlp_hf", exist_ok=True)
# Save the constructed weights
save_file({
"layer1.weight": W1,
"layer1.bias": b1,
"layer2.weight": W2,
"layer2.bias": b2,
}, "test_mlp_hf/model.safetensors")
print("\nSuccessfully saved constructed neuron to test_mlp_hf/model.safetensors") |