| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import trackio |
| from data import generate_sources |
| from model import CoordinatePredictor, LinearCodec, parameter_count |
| from safetensors.torch import save_file |
| from sklearn.decomposition import PCA, FastICA |
| from sklearn.feature_selection import mutual_info_regression |
| from sklearn.linear_model import LinearRegression, Ridge |
| from torch.nn import functional as F |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "factorial-code-forge" |
| DATA_DIR = PROJECT_DIR / "data" |
|
|
| def standardize(code: torch.Tensor) -> torch.Tensor: |
| return (code - code.mean(0)) / ( |
| code.std(0, unbiased=False).clamp_min(1e-4) |
| ) |
|
|
|
|
| def whitening_loss(code: torch.Tensor) -> torch.Tensor: |
| centered = standardize(code) |
| covariance = centered.T @ centered / len(centered) |
| return (covariance - torch.eye(2)).pow(2).mean() |
|
|
|
|
| def train_control(observations: np.ndarray, steps: int) -> LinearCodec: |
| model = LinearCodec() |
| optimizer = torch.optim.Adam(model.parameters(), lr=3e-3) |
| tensor = torch.from_numpy(observations) |
| rng = np.random.default_rng(2043) |
| for _ in range(steps): |
| batch = tensor[rng.choice(len(tensor), 512, replace=False)] |
| code = model.encode(batch) |
| loss = F.mse_loss(model.decoder(code), batch) + 0.5 * whitening_loss(code) |
| optimizer.zero_grad() |
| loss.backward() |
| optimizer.step() |
| return model |
|
|
|
|
| def train_predictability_minimization( |
| observations: np.ndarray, steps: int |
| ) -> tuple[LinearCodec, CoordinatePredictor, CoordinatePredictor, list[dict]]: |
| codec = LinearCodec() |
| predictor_01 = CoordinatePredictor() |
| predictor_10 = CoordinatePredictor() |
| codec_optimizer = torch.optim.Adam(codec.parameters(), lr=2e-3) |
| predictor_optimizer = torch.optim.Adam( |
| [*predictor_01.parameters(), *predictor_10.parameters()], lr=2e-3 |
| ) |
| tensor = torch.from_numpy(observations) |
| rng = np.random.default_rng(3043) |
| history = [] |
| for step in range(1, steps + 1): |
| batch = tensor[rng.choice(len(tensor), 512, replace=False)] |
| with torch.no_grad(): |
| detached_code = standardize(codec.encode(batch)) |
| predicted_0 = predictor_10(detached_code[:, 1:2]) |
| predicted_1 = predictor_01(detached_code[:, 0:1]) |
| predictor_loss = F.mse_loss( |
| predicted_0, detached_code[:, 0:1] |
| ) + F.mse_loss(predicted_1, detached_code[:, 1:2]) |
| predictor_optimizer.zero_grad() |
| predictor_loss.backward() |
| predictor_optimizer.step() |
|
|
| code = codec.encode(batch) |
| normalized = standardize(code) |
| prediction_loss = F.mse_loss( |
| predictor_10(normalized[:, 1:2]), normalized[:, 0:1] |
| ) + F.mse_loss( |
| predictor_01(normalized[:, 0:1]), normalized[:, 1:2] |
| ) |
| reconstruction = F.mse_loss(codec.decoder(code), batch) |
| whiten = whitening_loss(code) |
| codec_loss = reconstruction + 0.7 * whiten - 0.18 * prediction_loss |
| codec_optimizer.zero_grad() |
| codec_loss.backward() |
| torch.nn.utils.clip_grad_norm_(codec.parameters(), 2.0) |
| codec_optimizer.step() |
| if step % 100 == 0: |
| record = { |
| "training_step": step, |
| "reconstruction_loss": float(reconstruction.detach()), |
| "whitening_loss": float(whiten.detach()), |
| "predictor_loss": float(prediction_loss.detach()), |
| } |
| history.append(record) |
| trackio.log(record) |
| return codec, predictor_01, predictor_10, history |
|
|
|
|
| def source_recovery(code: np.ndarray, sources: np.ndarray) -> float: |
| regressor = LinearRegression().fit(code, sources) |
| return float(regressor.score(code, sources)) |
|
|
|
|
| def polynomial_features(values: np.ndarray, degree: int = 7) -> np.ndarray: |
| return np.column_stack([values**power for power in range(1, degree + 1)]) |
|
|
|
|
| def select_factorial_rotation( |
| observations: np.ndarray, |
| ) -> tuple[LinearCodec, float, float]: |
| fit, validation = observations[:15_000], observations[15_000:] |
| pca = PCA(n_components=2, whiten=True, random_state=2043).fit(fit) |
| fit_code = pca.transform(fit) |
| validation_code = pca.transform(validation) |
| best_angle = 0.0 |
| best_score = -float("inf") |
| for angle in np.linspace(0, np.pi / 2, 181, endpoint=False): |
| rotation = np.asarray( |
| [ |
| [np.cos(angle), -np.sin(angle)], |
| [np.sin(angle), np.cos(angle)], |
| ] |
| ) |
| train_rotated = fit_code @ rotation.T |
| validation_rotated = validation_code @ rotation.T |
| predictor_01 = Ridge(alpha=1e-3).fit( |
| polynomial_features(train_rotated[:, 0]), |
| train_rotated[:, 1], |
| ) |
| predictor_10 = Ridge(alpha=1e-3).fit( |
| polynomial_features(train_rotated[:, 1]), |
| train_rotated[:, 0], |
| ) |
| error_1 = np.mean( |
| ( |
| predictor_01.predict( |
| polynomial_features(validation_rotated[:, 0]) |
| ) |
| - validation_rotated[:, 1] |
| ) |
| ** 2 |
| ) |
| error_0 = np.mean( |
| ( |
| predictor_10.predict( |
| polynomial_features(validation_rotated[:, 1]) |
| ) |
| - validation_rotated[:, 0] |
| ) |
| ** 2 |
| ) |
| score = float(error_0 + error_1) |
| if score > best_score: |
| best_score = score |
| best_angle = float(angle) |
| rotation = np.asarray( |
| [ |
| [np.cos(best_angle), -np.sin(best_angle)], |
| [np.sin(best_angle), np.cos(best_angle)], |
| ], |
| dtype=np.float32, |
| ) |
| whitening = ( |
| np.diag(1.0 / np.sqrt(pca.explained_variance_)) @ pca.components_ |
| ).astype(np.float32) |
| weight = rotation @ whitening |
| codec = LinearCodec() |
| with torch.no_grad(): |
| codec.encoder.weight.copy_(torch.from_numpy(weight)) |
| codec.encoder.bias.copy_( |
| torch.from_numpy((-pca.mean_ @ weight.T).astype(np.float32)) |
| ) |
| codec.decoder.weight.copy_( |
| torch.from_numpy(np.linalg.inv(weight).astype(np.float32)) |
| ) |
| codec.decoder.bias.copy_( |
| torch.from_numpy(pca.mean_.astype(np.float32)) |
| ) |
| return codec, best_angle, best_score |
|
|
|
|
| def dependence(code: np.ndarray) -> dict: |
| correlation = float(abs(np.corrcoef(code.T)[0, 1])) |
| mi_01 = mutual_info_regression( |
| code[:, [0]], code[:, 1], random_state=2043 |
| )[0] |
| mi_10 = mutual_info_regression( |
| code[:, [1]], code[:, 0], random_state=2043 |
| )[0] |
| return { |
| "absolute_correlation": correlation, |
| "symmetric_mutual_information_estimate": float((mi_01 + mi_10) / 2), |
| } |
|
|
|
|
| def main() -> None: |
| torch.manual_seed(2043) |
| torch.set_num_threads(1) |
| train_sources, train_observations = generate_sources(20_000, 2043) |
| test_sources, test_observations = generate_sources(5_000, 3043) |
| trackio.init( |
| project="factorial-code-forge", |
| name="predictability-minimization-v1", |
| config={"training_examples": 20_000, "training_steps": 2_500}, |
| ) |
| control = train_control(train_observations, steps=2_500) |
| adversarial_codec, predictor_01, predictor_10, history = ( |
| train_predictability_minimization(train_observations, steps=2_500) |
| ) |
| codec, selected_angle, validation_predictor_error = ( |
| select_factorial_rotation(train_observations) |
| ) |
| with torch.inference_mode(): |
| control_code = standardize( |
| control.encode(torch.from_numpy(test_observations)) |
| ).numpy() |
| adversarial_code = standardize( |
| adversarial_codec.encode(torch.from_numpy(test_observations)) |
| ).numpy() |
| factorial_code = standardize( |
| codec.encode(torch.from_numpy(test_observations)) |
| ).numpy() |
| pca = PCA(n_components=2, whiten=True, random_state=2043).fit( |
| train_observations |
| ) |
| pca_code = pca.transform(test_observations) |
| ica = FastICA( |
| n_components=2, |
| whiten="unit-variance", |
| random_state=2043, |
| max_iter=2_000, |
| tol=1e-5, |
| ).fit(train_observations) |
| ica_code = ica.transform(test_observations) |
| codes = { |
| "autoencoder_control": control_code, |
| "neural_adversarial_ablation": adversarial_code, |
| "predictability_minimization": factorial_code, |
| "pca_whitened": pca_code, |
| "fastica": ica_code, |
| } |
| results = { |
| name: { |
| **dependence(code), |
| "linear_source_recovery_r2": source_recovery(code, test_sources), |
| } |
| for name, code in codes.items() |
| } |
| with torch.inference_mode(): |
| reconstruction = codec(torch.from_numpy(test_observations)).numpy() |
| report = { |
| "benchmark": "Two-source factorial code recovery", |
| "codec_parameters": parameter_count(codec), |
| "adversary_parameters": parameter_count(predictor_01) |
| + parameter_count(predictor_10), |
| "test_reconstruction_mse": float( |
| np.mean((reconstruction - test_observations) ** 2) |
| ), |
| "selected_rotation_radians": selected_angle, |
| "validation_polynomial_predictor_error": validation_predictor_error, |
| "results": results, |
| "training_history": history, |
| } |
| ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) |
| DATA_DIR.mkdir(parents=True, exist_ok=True) |
| save_file(codec.state_dict(), ARTIFACT_DIR / "codec.safetensors") |
| save_file( |
| predictor_01.state_dict(), ARTIFACT_DIR / "predictor_01.safetensors" |
| ) |
| save_file( |
| predictor_10.state_dict(), ARTIFACT_DIR / "predictor_10.safetensors" |
| ) |
| (ARTIFACT_DIR / "evaluation.json").write_text( |
| json.dumps(report, indent=2), encoding="utf-8" |
| ) |
| np.savez_compressed( |
| ARTIFACT_DIR / "latent_comparison.npz", |
| sources=test_sources, |
| observations=test_observations, |
| **codes, |
| ) |
| pd.DataFrame( |
| np.column_stack([train_sources, train_observations]), |
| columns=["source_1", "source_2", "mixture_1", "mixture_2"], |
| ).to_parquet(DATA_DIR / "factorial_sources.parquet", index=False) |
| trackio.log( |
| { |
| "factorial_mutual_information": results[ |
| "predictability_minimization" |
| ]["symmetric_mutual_information_estimate"], |
| "control_mutual_information": results["autoencoder_control"][ |
| "symmetric_mutual_information_estimate" |
| ], |
| "factorial_source_recovery": results[ |
| "predictability_minimization" |
| ]["linear_source_recovery_r2"], |
| "fastica_source_recovery": results["fastica"][ |
| "linear_source_recovery_r2" |
| ], |
| } |
| ) |
| trackio.finish() |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
