"""P073: frozen, reconstruction-only embeddings with small-label controls."""
from collections.abc import Iterator
from functools import partial
import lightning.pytorch as lit
import numpy as np
import pyarrow as pa
import torch
from reporting import report
import relflow as rf
PROOF_ID = "P073"
BUDGET = 600
TRAIN_ROWS = 4096
TEST_ROWS = 2048
LABELS_PER_CLASS = 8
CODES = ("arch", "bay", "cove", "dune", "elm", "ford", "glen", "hill")
CENTERS = np.asarray([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]])
def sample(*, rows: int, seed: int, permuted: bool = False) -> tuple[list[dict], np.ndarray]:
rng = np.random.default_rng(seed)
labels = np.resize(np.arange(4), rows)
rng.shuffle(labels)
codes = []
for _ in range(2):
noisy = np.where(rng.random(rows) < 0.25, rng.integers(4, size=rows), labels)
codes.append(2 * noisy + rng.integers(2, size=rows))
geometry = rng.normal(0.0, 2.0, (rows, 16))
geometry[:, :2] = CENTERS[labels] + rng.normal(0.0, 0.8, (rows, 2))
first, second = codes
if permuted:
first = first[rng.permutation(rows)]
second = second[rng.permutation(rows)]
geometry = geometry[rng.permutation(rows)]
records = [
{"first": CODES[a], "second": CODES[b], "geometry": vector.tolist()}
for a, b, vector in zip(first, second, geometry, strict=True)
]
return records, labels
def records(*, rows: int, seed: int, permuted: bool = False) -> Iterator[dict]:
# Labels stay outside the training stream and the model's schema.
yield from sample(rows=rows, seed=seed, permuted=permuted)[0]
def build() -> rf.Model:
mask = rf.Mask(rate=0.35, reconstruct=True)
return rf.Model.xs(
batch_size=64,
embed=True,
first=rf.Enum(values=CODES, p_unavailable=0.0, mask=mask),
second=rf.Enum(values=CODES, p_unavailable=0.0, mask=mask),
geometry=rf.Vector(n_dim=16, mask=mask),
)
def embedding(model: rf.Model, rows: list[dict]) -> np.ndarray:
output = model.predict(pa.Table.from_pylist(rows))["representations"].to_pylist()
roots = [
sorted(
(record for record in row if record["path"] == [] and record["parents"] == []),
key=lambda record: record["output_index"],
)
for row in output
]
return np.asarray([[record["embedding"] for record in root] for root in roots], dtype=np.float64).reshape(
len(rows), -1
)
def raw_features(rows: list[dict]) -> np.ndarray:
identity = np.eye(len(CODES))
return np.asarray(
[
np.concatenate((identity[CODES.index(row["first"])], identity[CODES.index(row["second"])], row["geometry"]))
for row in rows
]
)
def classify(reference: np.ndarray, labels: np.ndarray, queries: np.ndarray) -> np.ndarray:
"""Fit four centroids using only the supplied small labeled reference set."""
centroids = np.stack([reference[labels == label].mean(axis=0) for label in range(4)])
return np.square(queries[:, None, :] - centroids[None, :, :]).sum(axis=-1).argmin(axis=1)
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
train, _ = sample(rows=TRAIN_ROWS, seed=seed * 100 + 1)
reference, reference_labels = sample(rows=4 * LABELS_PER_CLASS, seed=seed * 100 + 3)
test, test_labels = sample(rows=TEST_ROWS, seed=seed * 100 + 4)
probe_rng = np.random.default_rng(seed * 100 + 5)
shuffled_labels = [probe_rng.permutation(reference_labels) for _ in range(32)]
metrics = {
"train_rows": TRAIN_ROWS,
"validation_rows": 512,
"test_rows": TEST_ROWS,
"probe_labels_per_class": LABELS_PER_CLASS,
"chance_accuracy": 0.25,
"arms": {},
}
checks = {}
for name in ("untrained", "reconstruction", "permuted_views"):
lit.seed_everything(seed, workers=True)
model = build()
trained_steps = 0
if name != "untrained":
model.optimizer = rf.adamw(learning_rate=0.002)
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=TRAIN_ROWS, seed=seed * 100 + 1, permuted=name == "permuted_views"),
validate=partial(records, rows=512, seed=seed * 100 + 2, permuted=name == "permuted_views"),
seed=seed,
)
trainer = lit.Trainer(
accelerator=accelerator,
devices=1,
max_epochs=-1,
max_steps=BUDGET if steps is None else min(steps, BUDGET),
deterministic=True,
logger=False,
enable_checkpointing=False,
enable_progress_bar=False,
enable_model_summary=False,
num_sanity_val_steps=0,
)
trainer.fit(model, datamodule=data)
trained_steps = trainer.global_step
model.eval().requires_grad_(False)
parameters = {key: value.detach().cpu().clone() for key, value in model.named_parameters()}
reference_vectors = embedding(model, reference)
test_vectors = embedding(model, test)
guessed = classify(reference_vectors, reference_labels, test_vectors)
shuffled_accuracies = [
float(np.mean(classify(reference_vectors, labels, test_vectors) == test_labels))
for labels in shuffled_labels
]
accuracy = float(np.mean(guessed == test_labels))
metrics["arms"][name] = {
"steps": trained_steps,
"accuracy": accuracy,
"shuffled_probe_accuracy": float(np.mean(shuffled_accuracies)),
"shuffled_probe_repetitions": len(shuffled_labels),
"embedding_dimensions": test_vectors.shape[1],
"norm_max_error": float(np.max(np.abs(np.linalg.norm(test_vectors, axis=1) - 1.0))),
}
checks[f"{name}: export is finite and unit-normalized"] = (
bool(np.isfinite(test_vectors).all()) and metrics["arms"][name]["norm_max_error"] < 1e-5
)
checks[f"{name}: export preserves every frozen parameter"] = all(
torch.equal(parameters[key], value.detach().cpu()) for key, value in model.named_parameters()
)
raw_train = raw_features(train)
mean, scale = raw_train.mean(axis=0), np.maximum(raw_train.std(axis=0), 1e-8)
raw_prediction = classify(
(raw_features(reference) - mean) / scale, reference_labels, (raw_features(test) - mean) / scale
)
metrics["raw_accuracy"] = float(np.mean(raw_prediction == test_labels))
trained = metrics["arms"]["reconstruction"]["accuracy"]
metrics["gain_over_raw"] = trained - metrics["raw_accuracy"]
checks["Reconstruction embeddings reach 0.70 accuracy with eight labels per class"] = trained >= 0.70
for control in ("untrained", "permuted_views"):
gain = trained - metrics["arms"][control]["accuracy"]
metrics[f"gain_over_{control}"] = gain
checks[f"Reconstruction exceeds {control} accuracy by 0.05"] = gain >= 0.05
checks["Shuffled probe labels remove class accuracy"] = (
metrics["arms"]["reconstruction"]["shuffled_probe_accuracy"] < 0.35
)
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=7301)
