Publish Predictability-minimized factorial code and adversarial ablation

README.md

@@ -0,0 +1,48 @@
+---
+license: apache-2.0
+tags:
+- independent-component-analysis
+- adversarial-learning
+- representation-learning
+- factorial-codes
+- gradio
+---
+
+# Factorial Code Forge
+
+Factorial Code Forge mixes two statistically independent non-Gaussian sources and
+asks a tiny autoencoder to recover a useful two-dimensional code. In the experimental
+variant, two adversarial predictors try to infer each latent coordinate from the
+other while the encoder tries to make those predictions fail. Reconstruction and
+variance constraints prevent a constant-code shortcut.
+
+The historical anchor is Schmidhuber's 1992
+[Learning Factorial Codes by Predictability
+Minimization](https://doi.org/10.1162/neco.1992.4.6.863). The modern experiment
+compares predictability minimization with an architecture-matched reconstructive
+encoder, PCA whitening, and FastICA. It is a new controlled benchmark, not a
+reproduction of the original paper's implementation or datasets.
+
+## Verified results
+
+Models trained on 20,000 mixed samples and were measured once on 5,000 new samples:
+
+| Code | Absolute correlation | Symmetric MI estimate |
+| --- | ---: | ---: |
+| Autoencoder control | 0.02200 | 0.22345 |
+| Online neural-adversarial ablation | 0.01886 | 0.22678 |
+| Validation-selected predictability minimum | 0.01158 | 0.00280 |
+| PCA whitening | 0.02039 | 0.22275 |
+| FastICA | 0.00590 | 0.00326 |
+
+The selected 12-parameter linear codec reconstructed the mixtures at
+`1.08e-14` MSE. Its rotation was selected using polynomial cross-predictors on a
+held-out portion of the training set, never the test set. The online neural
+minimax game did not improve nonlinear dependence and is retained as a negative
+stability result rather than omitted.
+
+## Reproduce
+
+```powershell
+uv run python projects/factorial-code-forge/train.py
+```

codec.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee01726e068b8d357d49d71b28916453915ad04b41643f1d32204cae1326d63d
+size 328

evaluation.json

@@ -0,0 +1,187 @@
+{
+  "benchmark": "Two-source factorial code recovery",
+  "codec_parameters": 12,
+  "adversary_parameters": 1346,
+  "test_reconstruction_mse": 1.0849382441166684e-14,
+  "selected_rotation_radians": 0.6335255903647925,
+  "validation_polynomial_predictor_error": 2.0412044050069174,
+  "results": {
+    "autoencoder_control": {
+      "absolute_correlation": 0.02199843078760486,
+      "symmetric_mutual_information_estimate": 0.22345146065238453,
+      "linear_source_recovery_r2": 1.0
+    },
+    "neural_adversarial_ablation": {
+      "absolute_correlation": 0.018864639802400177,
+      "symmetric_mutual_information_estimate": 0.22677681116066317,
+      "linear_source_recovery_r2": 1.0
+    },
+    "predictability_minimization": {
+      "absolute_correlation": 0.011581345795229356,
+      "symmetric_mutual_information_estimate": 0.002795710068105528,
+      "linear_source_recovery_r2": 1.0
+    },
+    "pca_whitened": {
+      "absolute_correlation": 0.02038519159561663,
+      "symmetric_mutual_information_estimate": 0.22274881647833844,
+      "linear_source_recovery_r2": 1.0
+    },
+    "fastica": {
+      "absolute_correlation": 0.005903947226495519,
+      "symmetric_mutual_information_estimate": 0.003260830744820531,
+      "linear_source_recovery_r2": 1.0
+    }
+  },
+  "training_history": [
+    {
+      "step": 100,
+      "reconstruction_loss": 0.5751938223838806,
+      "whitening_loss": 0.4999920129776001,
+      "predictor_loss": 0.03917834907770157
+    },
+    {
+      "step": 200,
+      "reconstruction_loss": 0.10436813533306122,
+      "whitening_loss": 0.4999969005584717,
+      "predictor_loss": 0.012760449200868607
+    },
+    {
+      "step": 300,
+      "reconstruction_loss": 0.09469099342823029,
+      "whitening_loss": 0.49983352422714233,
+      "predictor_loss": 0.0026066915597766638
+    },
+    {
+      "step": 400,
+      "reconstruction_loss": 0.0830593854188919,
+      "whitening_loss": 0.49922531843185425,
+      "predictor_loss": 0.004627525806427002
+    },
+    {
+      "step": 500,
+      "reconstruction_loss": 0.08060496300458908,
+      "whitening_loss": 0.4975007474422455,
+      "predictor_loss": 0.01004670187830925
+    },
+    {
+      "step": 600,
+      "reconstruction_loss": 0.07426909357309341,
+      "whitening_loss": 0.4930780827999115,
+      "predictor_loss": 0.02585626393556595
+    },
+    {
+      "step": 700,
+      "reconstruction_loss": 0.0617077499628067,
+      "whitening_loss": 0.48229295015335083,
+      "predictor_loss": 0.06387390196323395
+    },
+    {
+      "step": 800,
+      "reconstruction_loss": 0.040828146040439606,
+      "whitening_loss": 0.4678165912628174,
+      "predictor_loss": 0.1077573299407959
+    },
+    {
+      "step": 900,
+      "reconstruction_loss": 0.022691549733281136,
+      "whitening_loss": 0.3986736834049225,
+      "predictor_loss": 0.36198264360427856
+    },
+    {
+      "step": 1000,
+      "reconstruction_loss": 0.021030452102422714,
+      "whitening_loss": 0.23594792187213898,
+      "predictor_loss": 0.9913151860237122
+    },
+    {
+      "step": 1100,
+      "reconstruction_loss": 0.04022720083594322,
+      "whitening_loss": 0.009166543371975422,
+      "predictor_loss": 1.8695604801177979
+    },
+    {
+      "step": 1200,
+      "reconstruction_loss": 0.003848953638225794,
+      "whitening_loss": 0.00017961427511181682,
+      "predictor_loss": 1.7930779457092285
+    },
+    {
+      "step": 1300,
+      "reconstruction_loss": 0.0014683017507195473,
+      "whitening_loss": 0.004436715971678495,
+      "predictor_loss": 1.7725040912628174
+    },
+    {
+      "step": 1400,
+      "reconstruction_loss": 0.0005974461091682315,
+      "whitening_loss": 0.0012278903741389513,
+      "predictor_loss": 1.8158133029937744
+    },
+    {
+      "step": 1500,
+      "reconstruction_loss": 0.0003378715191502124,
+      "whitening_loss": 0.003598805284127593,
+      "predictor_loss": 1.7867335081100464
+    },
+    {
+      "step": 1600,
+      "reconstruction_loss": 6.85915001668036e-05,
+      "whitening_loss": 0.0005436374340206385,
+      "predictor_loss": 1.8233861923217773
+    },
+    {
+      "step": 1700,
+      "reconstruction_loss": 0.0001249544438906014,
+      "whitening_loss": 0.00044188593165017664,
+      "predictor_loss": 1.8262327909469604
+    },
+    {
+      "step": 1800,
+      "reconstruction_loss": 1.3117482012603432e-05,
+      "whitening_loss": 0.00011215780978091061,
+      "predictor_loss": 1.8308453559875488
+    },
+    {
+      "step": 1900,
+      "reconstruction_loss": 1.1593769158935174e-05,
+      "whitening_loss": 0.0016796954441815615,
+      "predictor_loss": 1.7771589756011963
+    },
+    {
+      "step": 2000,
+      "reconstruction_loss": 0.0004471217398531735,
+      "whitening_loss": 0.004560444969683886,
+      "predictor_loss": 1.7743375301361084
+    },
+    {
+      "step": 2100,
+      "reconstruction_loss": 2.6598576368996873e-05,
+      "whitening_loss": 0.0024056474212557077,
+      "predictor_loss": 1.7827980518341064
+    },
+    {
+      "step": 2200,
+      "reconstruction_loss": 4.374821946839802e-05,
+      "whitening_loss": 7.564837142126635e-05,
+      "predictor_loss": 1.8309648036956787
+    },
+    {
+      "step": 2300,
+      "reconstruction_loss": 0.00020850264991167933,
+      "whitening_loss": 0.005683479364961386,
+      "predictor_loss": 1.822058916091919
+    },
+    {
+      "step": 2400,
+      "reconstruction_loss": 2.0745503206853755e-05,
+      "whitening_loss": 8.849342520989012e-07,
+      "predictor_loss": 1.8361434936523438
+    },
+    {
+      "step": 2500,
+      "reconstruction_loss": 1.0664540241123177e-05,
+      "whitening_loss": 0.0009145042859017849,
+      "predictor_loss": 1.7433021068572998
+    }
+  ]
+}

latent_comparison.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efcdab9ba10b29127fb11ed21ae13df08e1b443ce6becc0f018163495cd7e050
+size 260780

predictor_01.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f637d7931c2ae64658844975ee2107a5a5a3629af5a9d063b22a1dcef91a67f
+size 3140

predictor_10.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d4ffe8417e116e9a2eb3c0dd63cf3e7c5c9fbfd68b33859b11aa041b9088501
+size 3140

source/app.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+import plotly.graph_objects as go
+
+ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "factorial-code-forge"
+DATA = np.load(ARTIFACT_DIR / "latent_comparison.npz")
+METHODS = {
+    "Observed mixtures": "observations",
+    "Autoencoder control": "autoencoder_control",
+    "Neural adversarial ablation": "neural_adversarial_ablation",
+    "Predictability minimization": "predictability_minimization",
+    "PCA whitened": "pca_whitened",
+    "FastICA": "fastica",
+}
+
+
+def explore(method: str, points: int) -> tuple[go.Figure, dict]:
+    key = METHODS[method]
+    code = DATA[key][: int(points)]
+    sources = DATA["sources"][: int(points)]
+    figure = go.Figure(
+        go.Scattergl(
+            x=code[:, 0],
+            y=code[:, 1],
+            mode="markers",
+            marker={
+                "size": 4,
+                "color": sources[:, 0],
+                "colorscale": "Turbo",
+                "opacity": 0.65,
+            },
+        )
+    )
+    figure.update_layout(
+        title=f"{method}: learned two-dimensional code",
+        xaxis_title="Coordinate 1",
+        yaxis_title="Coordinate 2",
+        template="plotly_dark",
+    )
+    return figure, {
+        "absolute_correlation": round(float(abs(np.corrcoef(code.T)[0, 1])), 5),
+        "points": len(code),
+        "color": "true source 1 (evaluation only)",
+    }
+
+
+with gr.Blocks(title="Factorial Code Forge") as demo:
+    gr.Markdown(
+        "# Factorial Code Forge\n"
+        "Inspect how different unsupervised methods reorganize two mixed, "
+        "independent non-Gaussian sources."
+    )
+    with gr.Row():
+        method = gr.Dropdown(
+            list(METHODS), value="Predictability minimization", label="Code"
+        )
+        points = gr.Slider(250, 5_000, 2_000, step=250, label="Points")
+    run = gr.Button("Reveal code", variant="primary")
+    scatter = gr.Plot()
+    metrics = gr.JSON()
+    run.click(explore, [method, points], [scatter, metrics])
+    demo.load(explore, [method, points], [scatter, metrics])
+
+
+if __name__ == "__main__":
+    demo.launch()

source/data.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def generate_sources(samples: int, seed: int) -> tuple[np.ndarray, np.ndarray]:
+    rng = np.random.default_rng(seed)
+    source1 = rng.laplace(size=samples) / np.sqrt(2)
+    source2 = rng.uniform(-np.sqrt(3), np.sqrt(3), size=samples)
+    sources = np.column_stack([source1, source2]).astype(np.float32)
+    mixing = np.asarray([[1.0, 0.85], [0.55, 1.25]], dtype=np.float32)
+    observations = sources @ mixing.T
+    return sources, observations.astype(np.float32)

source/model.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class LinearCodec(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.encoder = nn.Linear(2, 2)
+        self.decoder = nn.Linear(2, 2)
+
+    def encode(self, observations: torch.Tensor) -> torch.Tensor:
+        return self.encoder(observations)
+
+    def forward(self, observations: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.encode(observations))
+
+
+class CoordinatePredictor(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.network = nn.Sequential(
+            nn.Linear(1, 24),
+            nn.Tanh(),
+            nn.Linear(24, 24),
+            nn.Tanh(),
+            nn.Linear(24, 1),
+        )
+
+    def forward(self, coordinate: torch.Tensor) -> torch.Tensor:
+        return self.network(coordinate)
+
+
+def parameter_count(module: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in module.parameters())

source/requirements.txt

@@ -0,0 +1,9 @@
+gradio>=5,<7
+numpy>=2,<3
+pandas>=2.3,<4
+plotly>=6,<7
+pyarrow>=21,<24
+safetensors>=0.6,<1
+scikit-learn>=1.7,<2
+torch>=2.7,<3
+trackio>=0.16,<1

source/train.py

@@ -0,0 +1,309 @@
+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()