"""P054: leakage-free prefix conditioning and free-running arithmetic-sequence rollout."""
from collections.abc import Iterator
from functools import partial
import lightning.pytorch as lit
import numpy as np
import pyarrow as pa
from reporting import report
import relflow as rf
PROOF_ID = "P054"
LENGTH = 5
BUDGET = 800
def records(*, rows: int, seed: int) -> Iterator[dict]:
rng = np.random.default_rng(seed)
for _ in range(rows):
offset = float(rng.uniform(-1.0, 1.0))
increment = float(rng.uniform(-0.4, 0.4))
cut = int(rng.integers(2, LENGTH))
yield {
"events": [
{"value": offset + index * increment, "future": index >= cut, "next": index == cut}
for index in range(LENGTH)
]
}
def build() -> rf.Model:
return rf.Model(
d_model=32,
n_layers=2,
n_heads=4,
dropout=0.0,
batch_size=64,
events=rf.Branch(
length=LENGTH,
reduction=None,
mask=rf.Mask(query="future", skip=True, dropout=False),
value=rf.Number(objective="mse", mask=rf.Mask(query="next", reconstruct=True)),
),
)
def prediction(model: rf.Model, rows: list[dict]) -> tuple[np.ndarray, np.ndarray]:
output = model.predict(pa.Table.from_pylist(rows))["predictions"].to_pylist()
content = np.asarray(
[[item["value"]["content"] for item in row["events"]] for row in output], dtype=np.float64
).reshape(len(rows), LENGTH)
inferred = np.asarray(
[[item["value"]["inferred"] for item in row["events"]] for row in output], dtype=bool
).reshape(len(rows), LENGTH)
return content, inferred
def rollout(model: rf.Model, rows: list[dict]) -> np.ndarray:
"""Expose generated values only; never teacher-force the held-out suffix."""
values = np.zeros((len(rows), LENGTH), dtype=np.float64)
values[:, :2] = [[event["value"] for event in row["events"][:2]] for row in rows]
for cut in range(2, LENGTH):
inputs = [
{
"events": [
{"value": float(value), "future": index >= cut, "next": index == cut}
for index, value in enumerate(sequence)
]
}
for sequence in values
]
content, _ = prediction(model, inputs)
values[:, cut] = content[:, cut]
return values
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
lit.seed_everything(seed, workers=True)
model = build()
model.optimizer = rf.adamw(learning_rate=0.002)
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=4096, seed=seed + 1),
validate=partial(records, rows=512, 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),
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)
model.eval()
train = list(records(rows=4096, seed=seed + 1))
test = list(records(rows=1024, seed=seed + 3))
actual = np.asarray([[event["value"] for event in row["events"]] for row in test])
mean = np.asarray([[event["value"] for event in row["events"]] for row in train]).mean(axis=0)
order = np.random.default_rng(seed + 4).permutation(len(test))
metrics = {"steps": trainer.global_step, "prefixes": {}}
checks = {}
# A separate model audits target selection without fitting held-out moments.
probe = build()
for cut in range(2, LENGTH):
selected = [
{
"events": [
{**event, "future": index >= cut, "next": index == cut} for index, event in enumerate(row["events"])
]
}
for row in test
]
content, inferred = prediction(model, selected)
poisoned = [
{
"events": [
{**event, "value": 1000.0 + index} if index >= cut else event
for index, event in enumerate(row["events"])
]
}
for row in selected
]
shuffled = [
{
"events": [
{**event, "value": float(actual[other, index])} if index < cut else event
for index, event in enumerate(row["events"])
]
}
for row, other in zip(selected, order, strict=True)
]
replaced, _ = prediction(model, poisoned)
broken, _ = prediction(model, shuffled)
baseline = float(np.sqrt(np.mean((actual[:, cut] - mean[cut]) ** 2)))
measured = {
"nrmse": float(np.sqrt(np.mean((content[:, cut] - actual[:, cut]) ** 2))) / baseline,
"baseline_rmse": baseline,
"shuffled_nrmse": float(np.sqrt(np.mean((broken[:, cut] - actual[:, cut]) ** 2))) / baseline,
"future_value_drift": float(np.max(np.abs(content[:, cut] - replaced[:, cut]))),
}
metrics["prefixes"][cut] = measured
checks[f"prefix {cut}: next-value nRMSE below 0.35"] = measured["nrmse"] < 0.35
checks[f"prefix {cut}: shuffled prefix nRMSE above 0.85"] = measured["shuffled_nrmse"] > 0.85
checks[f"prefix {cut}: future cannot affect prediction"] = measured["future_value_drift"] < 1e-5
field = probe.encode(pa.Table.from_pylist(selected), strata="train")["/events/value"]
expected = np.broadcast_to(np.arange(LENGTH) == cut, inferred.shape)
visible = np.broadcast_to(np.arange(LENGTH) < cut, inferred.shape)
checks[f"prefix {cut}: exactly the next coordinate is requested"] = bool(np.array_equal(inferred, expected))
checks[f"prefix {cut}: exactly the next coordinate is trained"] = bool(
np.array_equal(field.trainable.cpu().numpy().reshape(expected.shape), expected)
)
checks[f"prefix {cut}: only prefix values are encoder inputs"] = bool(
np.array_equal(field.present.cpu().numpy().reshape(visible.shape), visible)
)
generated = rollout(model, test)
baseline = float(np.sqrt(np.mean((actual[:, 2:] - mean[2:]) ** 2)))
rmse = float(np.sqrt(np.mean((generated[:, 2:] - actual[:, 2:]) ** 2)))
persistence = float(np.sqrt(np.mean((actual[:, 1:2] - actual[:, 2:]) ** 2)))
metrics["rollout"] = {
"rmse": rmse,
"baseline_rmse": baseline,
"nrmse": rmse / baseline,
"persistence_rmse": persistence,
}
checks["Free-running suffix nRMSE below 0.45"] = metrics["rollout"]["nrmse"] < 0.45
checks["Rollout beats persistence"] = rmse < persistence
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=5401)
