Publish Continuous-decay recurrent model and matched controls

README.md

@@ -0,0 +1,33 @@
+---
+title: Liquid Time Pocket
+emoji: 💧
+colorFrom: cyan
+colorTo: blue
+sdk: gradio
+sdk_version: "6.5.1"
+app_file: app.py
+pinned: false
+---
+
+# Liquid Time Pocket
+
+Liquid Time Pocket trains a recurrent cell whose hidden state relaxes toward a
+learned candidate through per-unit continuous time constants. The GRU and plain
+RNN controls contain exactly the same 1,887 parameters and receive delta time as
+an ordinary input feature.
+
+All models train with gaps from 0.02 to 0.12, then face unseen gaps up to 0.40 on
+twice-longer sequences. The Space plots predictions against continuous timestamps.
+
+## Verified local result
+
+At exactly 1,887 parameters, normal-gap RMSE was tightly matched: 0.0893 for the
+liquid cell, 0.0869 for the GRU, and 0.0881 for the RNN. Under unseen gaps up to
+0.40, the liquid cell failed at 1.514 RMSE versus 0.274/0.281 for GRU/RNN. Its
+learned time constants ranged from 0.122 to 0.716, causing excessive phase-memory
+decay under the shifted interval distribution.
+
+```bash
+uv run python projects/liquid-time-pocket/train.py
+uv run pytest tests/test_liquid_time_pocket.py
+```

evaluation.json

@@ -0,0 +1,82 @@
+{
+  "experiment": "Liquid time-constant irregular forecasting",
+  "learned_time_constants": {
+    "minimum": 0.12157730013132095,
+    "mean": 0.45884114503860474,
+    "maximum": 0.7155194282531738
+  },
+  "results": {
+    "liquid_time_constant": {
+      "parameters": 1887,
+      "best_step": 2500,
+      "normal_gaps": {
+        "rmse": 0.08929167687892914,
+        "mae": 0.03883583843708038,
+        "examples": 2000,
+        "sequence_length": 128,
+        "delta_time_range": [
+          0.02,
+          0.12
+        ]
+      },
+      "unseen_large_gaps": {
+        "rmse": 1.5136159658432007,
+        "mae": 1.1862934827804565,
+        "examples": 2000,
+        "sequence_length": 128,
+        "delta_time_range": [
+          0.12,
+          0.4
+        ]
+      }
+    },
+    "matched_gru": {
+      "parameters": 1887,
+      "best_step": 2500,
+      "normal_gaps": {
+        "rmse": 0.08689922839403152,
+        "mae": 0.040582284331321716,
+        "examples": 2000,
+        "sequence_length": 128,
+        "delta_time_range": [
+          0.02,
+          0.12
+        ]
+      },
+      "unseen_large_gaps": {
+        "rmse": 0.27418622374534607,
+        "mae": 0.21493969857692719,
+        "examples": 2000,
+        "sequence_length": 128,
+        "delta_time_range": [
+          0.12,
+          0.4
+        ]
+      }
+    },
+    "matched_rnn": {
+      "parameters": 1887,
+      "best_step": 2500,
+      "normal_gaps": {
+        "rmse": 0.08811825513839722,
+        "mae": 0.04227612167596817,
+        "examples": 2000,
+        "sequence_length": 128,
+        "delta_time_range": [
+          0.02,
+          0.12
+        ]
+      },
+      "unseen_large_gaps": {
+        "rmse": 0.2811814546585083,
+        "mae": 0.21727129817008972,
+        "examples": 2000,
+        "sequence_length": 128,
+        "delta_time_range": [
+          0.12,
+          0.4
+        ]
+      }
+    }
+  }
+}

liquid_time_constant.safetensors

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

matched_gru.safetensors

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

matched_rnn.safetensors

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

source/app.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import gradio as gr
+import plotly.graph_objects as go
+import torch
+from data import irregular_batch
+from model import LiquidTimeConstantRNN, MatchedGRU, MatchedRNN
+from safetensors.torch import load_file
+
+ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "liquid-time-pocket"
+REPORT = json.loads((ARTIFACT_DIR / "evaluation.json").read_text(encoding="utf-8"))
+MODELS = {
+    "Liquid time constant": (LiquidTimeConstantRNN(), "liquid_time_constant"),
+    "Matched GRU": (MatchedGRU(), "matched_gru"),
+    "Matched RNN": (MatchedRNN(), "matched_rnn"),
+}
+for model, key in MODELS.values():
+    model.load_state_dict(load_file(ARTIFACT_DIR / f"{key}.safetensors"))
+    model.eval()
+
+
+@torch.inference_mode()
+def forecast(seed: int, large_gaps: bool) -> tuple[go.Figure, dict]:
+    inputs, targets, time = irregular_batch(
+        1, 128, int(seed), large_gaps=bool(large_gaps)
+    )
+    figure = go.Figure()
+    figure.add_trace(
+        go.Scatter(x=time[0], y=targets[0, :, 0], name="Target", line={"width": 4})
+    )
+    metrics = {}
+    for label, (model, key) in MODELS.items():
+        prediction = model(inputs)[0, :, 0]
+        figure.add_trace(go.Scatter(x=time[0], y=prediction, name=label))
+        metrics[label] = {
+            "live_rmse": float(
+                (prediction - targets[0, :, 0]).square().mean().sqrt()
+            ),
+            "verified_large_gap_rmse": REPORT["results"][key][
+                "unseen_large_gaps"
+            ]["rmse"],
+        }
+    figure.update_layout(
+        template="plotly_dark",
+        title="Irregularly sampled multiscale forecast",
+        xaxis_title="Continuous time",
+        yaxis_title="Signal",
+    )
+    return figure, metrics
+
+
+with gr.Blocks(title="Liquid Time Pocket") as demo:
+    gr.Markdown(
+        "# Liquid Time Pocket\n"
+        "Compare a learned continuous decay cell with exactly parameter-matched "
+        "GRU and RNN controls under irregular time gaps."
+    )
+    with gr.Row():
+        seed = gr.Slider(0, 100_000, value=2203, step=1, label="Task seed")
+        large_gaps = gr.Checkbox(True, label="Use unseen large time gaps")
+    initial = forecast(2203, True)
+    chart = gr.Plot(value=initial[0])
+    metrics = gr.JSON(value=initial[1])
+    button = gr.Button("Forecast irregular telemetry", variant="primary")
+    button.click(forecast, inputs=[seed, large_gaps], outputs=[chart, metrics])
+
+
+if __name__ == "__main__":
+    demo.launch()
+

source/data.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+
+def irregular_batch(
+    batch: int,
+    length: int,
+    seed: int,
+    *,
+    large_gaps: bool,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    rng = np.random.default_rng(seed)
+    low, high = ((0.12, 0.40) if large_gaps else (0.02, 0.12))
+    delta_time = rng.uniform(low, high, size=(batch, length))
+    time = np.cumsum(delta_time, axis=1)
+    amplitudes = rng.uniform(0.5, 1.5, size=(batch, 3, 1))
+    frequencies = rng.uniform(
+        np.array([0.7, 1.7, 3.5])[None, :, None],
+        np.array([1.1, 2.3, 5.0])[None, :, None],
+        size=(batch, 3, 1),
+    )
+    phases = rng.uniform(0, 2 * np.pi, size=(batch, 3, 1))
+    values = (
+        amplitudes
+        * np.sin(frequencies * time[:, None, :] + phases)
+        * np.array([1.0, 0.5, 0.2])[None, :, None]
+    ).sum(axis=1)
+    values += rng.normal(0, 0.01, values.shape)
+    current = np.concatenate([np.zeros((batch, 1)), values[:, :-1]], axis=1)
+    inputs = np.stack([current, delta_time], axis=2).astype(np.float32)
+    return (
+        torch.from_numpy(inputs),
+        torch.from_numpy(values.astype(np.float32)).unsqueeze(2),
+        torch.from_numpy(time.astype(np.float32)),
+    )
+

source/model.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class LiquidTimeConstantRNN(nn.Module):
+    def __init__(self, hidden_dimensions: int = 41) -> None:
+        super().__init__()
+        self.hidden_dimensions = hidden_dimensions
+        self.candidate = nn.Linear(hidden_dimensions + 2, hidden_dimensions)
+        self.log_time_constant = nn.Parameter(torch.zeros(hidden_dimensions))
+        self.output = nn.Linear(hidden_dimensions, 1)
+
+    def forward(self, sequence: torch.Tensor) -> torch.Tensor:
+        hidden = torch.zeros(
+            len(sequence),
+            self.hidden_dimensions,
+            device=sequence.device,
+        )
+        time_constant = torch.nn.functional.softplus(self.log_time_constant) + 0.03
+        outputs = []
+        for step in range(sequence.shape[1]):
+            value_and_dt = sequence[:, step]
+            candidate_inputs = torch.cat(
+                [
+                    value_and_dt[:, :1],
+                    value_and_dt[:, 1:2].clamp(max=0.12),
+                ],
+                dim=1,
+            )
+            candidate = torch.tanh(
+                self.candidate(torch.cat([candidate_inputs, hidden], dim=1))
+            )
+            delta_time = value_and_dt[:, 1:2]
+            decay = torch.exp(-delta_time / time_constant[None])
+            hidden = decay * hidden + (1 - decay) * candidate
+            outputs.append(self.output(hidden))
+        return torch.stack(outputs, dim=1)
+
+
+class MatchedGRU(nn.Module):
+    def __init__(self, hidden_dimensions: int = 23) -> None:
+        super().__init__()
+        self.recurrent = nn.GRU(2, hidden_dimensions, batch_first=True)
+        self.output = nn.Linear(hidden_dimensions, 1)
+
+    def forward(self, sequence: torch.Tensor) -> torch.Tensor:
+        hidden, _ = self.recurrent(sequence)
+        return self.output(hidden)
+
+
+class MatchedRNN(nn.Module):
+    def __init__(self, hidden_dimensions: int = 41) -> None:
+        super().__init__()
+        self.recurrent = nn.RNN(2, hidden_dimensions, batch_first=True)
+        self.output = nn.Linear(hidden_dimensions, 1)
+
+    def forward(self, sequence: torch.Tensor) -> torch.Tensor:
+        hidden, _ = self.recurrent(sequence)
+        return self.output(hidden)
+
+
+def parameter_count(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters())

source/requirements.txt

@@ -0,0 +1,5 @@
+gradio>=5,<7
+plotly>=6,<7
+safetensors>=0.6,<1
+torch>=2.7,<3
+

source/train.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pandas as pd
+import torch
+import trackio
+from data import irregular_batch
+from model import (
+    LiquidTimeConstantRNN,
+    MatchedGRU,
+    MatchedRNN,
+    parameter_count,
+)
+from safetensors.torch import save_file
+from torch.nn import functional as F
+
+PROJECT_DIR = Path(__file__).resolve().parent
+ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "liquid-time-pocket"
+DATA_DIR = PROJECT_DIR / "data"
+SEED = 2203
+
+
+@torch.inference_mode()
+def evaluate(model: torch.nn.Module, *, large_gaps: bool) -> dict:
+    inputs, targets, _ = irregular_batch(
+        2_000,
+        128,
+        SEED + 10_000 + int(large_gaps),
+        large_gaps=large_gaps,
+    )
+    prediction = model(inputs)
+    error = prediction - targets
+    return {
+        "rmse": float(error.square().mean().sqrt()),
+        "mae": float(error.abs().mean()),
+        "examples": len(inputs),
+        "sequence_length": inputs.shape[1],
+        "delta_time_range": [0.12, 0.40] if large_gaps else [0.02, 0.12],
+    }
+
+
+def train_variant(name: str, model: torch.nn.Module) -> tuple[torch.nn.Module, int]:
+    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-5)
+    best = float("inf")
+    best_step = 0
+    best_state = None
+    for step in range(1, 2_501):
+        inputs, targets, _ = irregular_batch(
+            96, 64, SEED + step, large_gaps=False
+        )
+        prediction = model(inputs)
+        loss = F.mse_loss(prediction, targets)
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
+        optimizer.step()
+        if step % 125 == 0:
+            validation = evaluate(model, large_gaps=False)
+            trackio.log(
+                {
+                    "variant": name,
+                    "training_step": step,
+                    "training_mse": float(loss.detach()),
+                    "validation_rmse": validation["rmse"],
+                }
+            )
+            if validation["rmse"] < best:
+                best = validation["rmse"]
+                best_step = step
+                best_state = {
+                    key: value.detach().cpu().clone()
+                    for key, value in model.state_dict().items()
+                }
+    assert best_state is not None
+    model.load_state_dict(best_state)
+    return model, best_step
+
+
+def main() -> None:
+    torch.manual_seed(SEED)
+    torch.set_num_threads(1)
+    models = {
+        "liquid_time_constant": LiquidTimeConstantRNN(),
+        "matched_gru": MatchedGRU(),
+        "matched_rnn": MatchedRNN(),
+    }
+    assert {parameter_count(model) for model in models.values()} == {1_887}
+    trackio.init(
+        project="liquid-time-pocket",
+        name="irregular-gap-ltc-v1",
+        config={
+            "parameters_per_model": 1_887,
+            "training_steps": 2_500,
+            "training_delta_time": [0.02, 0.12],
+            "unseen_delta_time": [0.12, 0.40],
+        },
+    )
+    results = {}
+    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
+    DATA_DIR.mkdir(parents=True, exist_ok=True)
+    for name, model in models.items():
+        model, best_step = train_variant(name, model)
+        results[name] = {
+            "parameters": parameter_count(model),
+            "best_step": best_step,
+            "normal_gaps": evaluate(model, large_gaps=False),
+            "unseen_large_gaps": evaluate(model, large_gaps=True),
+        }
+        save_file(model.state_dict(), ARTIFACT_DIR / f"{name}.safetensors")
+    inputs, targets, time = irregular_batch(
+        100, 128, SEED + 30_000, large_gaps=True
+    )
+    frame = {
+        "task": [],
+        "time": [],
+        "target": [],
+        **{name: [] for name in models},
+    }
+    with torch.inference_mode():
+        predictions = {name: model(inputs) for name, model in models.items()}
+    for task in range(len(inputs)):
+        for step in range(inputs.shape[1]):
+            frame["task"].append(task)
+            frame["time"].append(float(time[task, step]))
+            frame["target"].append(float(targets[task, step, 0]))
+            for name in models:
+                frame[name].append(float(predictions[name][task, step, 0]))
+    report = {
+        "experiment": "Liquid time-constant irregular forecasting",
+        "learned_time_constants": {
+            "minimum": float(
+                (F.softplus(models["liquid_time_constant"].log_time_constant) + 0.03)
+                .min()
+            ),
+            "mean": float(
+                (F.softplus(models["liquid_time_constant"].log_time_constant) + 0.03)
+                .mean()
+            ),
+            "maximum": float(
+                (F.softplus(models["liquid_time_constant"].log_time_constant) + 0.03)
+                .max()
+            ),
+        },
+        "results": results,
+    }
+    (ARTIFACT_DIR / "evaluation.json").write_text(
+        json.dumps(report, indent=2), encoding="utf-8"
+    )
+    pd.DataFrame(frame).to_parquet(DATA_DIR / "large_gap_forecasts.parquet", index=False)
+    trackio.log(
+        {
+            f"{name}_large_gap_rmse": result["unseen_large_gaps"]["rmse"]
+            for name, result in results.items()
+        }
+    )
+    trackio.finish()
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()