"""P067: learn a skewed numeric relationship through Quantile representations."""
from collections.abc import Iterator
from functools import partial
from pathlib import Path
from tempfile import TemporaryDirectory
import lightning.pytorch as lit
import numpy as np
import pyarrow as pa
from reporting import report
import relflow as rf
PROOF_ID = "P067"
BUDGET = 400
BATCH_SIZE = 128
TRAIN_ROWS = 4096
TEST_ROWS = 2048
def records(*, rows: int, seed: int) -> Iterator[dict]:
rng = np.random.default_rng(seed)
amount = 50.0 * np.expm1(3.0 * rng.uniform(0.0, 1.0, rows))
cost = 200.0 + 3.0 * amount + rng.normal(0.0, 10.0, rows)
for value, target in zip(amount, cost, strict=True):
yield {"amount": float(value), "cost": float(target)}
def build() -> rf.Model:
model = rf.Model.xs(
batch_size=BATCH_SIZE,
amount=rf.Quantile,
cost=rf.Quantile(objective="mse", mask=True),
)
model.optimizer = rf.adamw(learning_rate=0.002)
return model
def predict(model: rf.Model, rows: list[dict]) -> np.ndarray:
predictions = model.predict(pa.Table.from_pylist(rows))["predictions"].to_pylist()
return np.asarray([row["cost"]["content"] for row in predictions], dtype=np.float64)
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
lit.seed_everything(seed, workers=True)
model = build()
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=TRAIN_ROWS, seed=seed + 1),
validate=partial(records, rows=1024, seed=seed + 2),
seed=seed,
)
trainer = lit.Trainer(
accelerator=accelerator,
devices=1,
max_epochs=-1,
max_steps=BUDGET if steps is None else min(steps, BUDGET),
logger=False,
enable_checkpointing=False,
enable_progress_bar=False,
enable_model_summary=False,
deterministic=True,
num_sanity_val_steps=0,
)
trainer.fit(model, datamodule=data)
model.eval()
addresses = (rf.Address("amount"), rf.Address("cost"))
normalization = {str(address): rf.Quantile.normalization(model, address) for address in addresses}
rows = list(records(rows=TEST_ROWS, seed=seed + 3))
targets = np.asarray([row["cost"] for row in rows])
visible = [{"amount": row["amount"]} for row in rows]
training_mean = float(np.mean([row["cost"] for row in records(rows=TRAIN_ROWS, seed=seed + 1)]))
baseline_rmse = float(np.sqrt(np.mean((training_mean - targets) ** 2)))
output = predict(model, visible)
shuffled = np.random.default_rng(seed + 4).permutation(TEST_ROWS)
control = predict(model, [visible[index] for index in shuffled])
rmse = float(np.sqrt(np.mean((output - targets) ** 2)))
control_rmse = float(np.sqrt(np.mean((control - targets) ** 2)))
original = predict(model, rows[:128])
tampered = predict(model, [{**row, "cost": row["cost"] + 1e6} for row in rows[:128]])
target_delta = float(np.max(np.abs(original - tampered)))
evaluation_frozen = normalization == {
str(address): rf.Quantile.normalization(model, address) for address in addresses
}
with TemporaryDirectory(prefix="relflow-quantile-proof-") as directory:
path = Path(directory) / "model.pt"
model.save(path)
restored = rf.Model.load(path).to(model.device)
restored.eval()
checkpoint_delta = float(np.max(np.abs(original - predict(restored, rows[:128]))))
checkpoint_state = normalization == {
str(address): rf.Quantile.normalization(restored, address) for address in addresses
}
metrics = {
"steps": trainer.global_step,
"train_rows": TRAIN_ROWS,
"test_rows": TEST_ROWS,
"training_mean": training_mean,
"baseline_rmse": baseline_rmse,
"noise_rmse": 10.0,
"rmse": rmse,
"normalized_rmse": rmse / baseline_rmse,
"shuffled_rmse": control_rmse,
"shuffled_normalized_rmse": control_rmse / baseline_rmse,
"normalization": normalization,
"target_tampering_max_delta": target_delta,
"checkpoint_max_delta": checkpoint_delta,
}
checks = {
"Held-out source-unit normalized RMSE below 0.25": rmse / baseline_rmse < 0.25,
"Shuffled input normalized RMSE above 0.90": control_rmse / baseline_rmse > 0.90,
"Shuffling more than triples the RMSE": control_rmse > 3.0 * rmse,
"Predictions in original units are finite": bool(np.isfinite(output).all() and np.isfinite(control).all()),
"Both digests count exactly the consumed training rows": all(
state["count"] == trainer.global_step * BATCH_SIZE for state in normalization.values()
),
"Evaluation preserves both training distributions": evaluation_frozen,
"Hidden target changes do not affect predictions": target_delta < 1e-8,
"Checkpoint preserves both training distributions": checkpoint_state,
"Checkpoint predictions remain within 0.0001 source units": checkpoint_delta < 1e-4,
}
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=6701)
