GenODE Inference

GenODE learns inference clocks for frozen generative backbones. A shared GICO teacher supervises either a deterministic density policy or an autoregressive stochastic density policy. Each task/backbone has separately trained weights. The generative backbone stays frozen.

Install

Python 3.11 or newer is required.

git clone https://github.com/pixelhero98/GenODE-Inference.git
cd GenODE-Inference
python -m pip install -e ".[test]"
python -m pytest -q

The optional latent-clock extra supplies Bayesian-optimization dependencies. Image generators, text-to-image scorers, and pretrained weights are external assets; supply their source revisions and checkpoints explicitly. See THIRD_PARTY_NOTICES.md.

Tasks and rewards

Task	Frozen native context	Terminal training reward
solar_energy_10m, traffic_hourly, weather_daily	Pooled backbone summary, including native auxiliary conditioning	Equal-weight CRPS and MASE log improvements
molecule_3d_set1, molecule_3d_set2, molecule_3d_set3	Pooled backbone summary of observed history	40% Kabsch RMSD + 15% each of four motion-discrepancy log improvements
cifar10	Explicit zero vector	Uniform KID minus candidate KID on paired generated/real-reference blocks
imagenet64	Native class embedding	Equally weighted class-conditional paired KID improvement
sana, sd15	Pooled native text embedding	Equal-weight ImageReward and VQAScore differences, divided by frozen pilot component scales

Pair each candidate and uniform anchor on context, backbone, solver, NFE, generation seed, ensemble size, reference data, and measurement protocol. Standard image GICO uses complete generated/real-reference feature blocks, paired across candidate and uniform clocks. ImageNet blocks are class conditional. Average repeated terminal measurements within each comparison cell before computing improvements; log ratios apply only to the positive-error tasks below.

Positive errors use log((anchor + epsilon) / (candidate + epsilon)), where each frozen numerical floor is 1e-6 * median(positive uniform calibration values). Reject degenerate calibration. After component scalarization, divide by a single frozen reward standard deviation for each task/backbone/solver, balancing calibration contributions across training NFEs. Do not subtract a mean or use per-context or running normalization. Uniform rewards remain exactly zero. Text-to-image component scales use pilot measurements; its scalar scale uses training measurements.

Image GICO uses KID supervision on RF++ and the other native image backbones (EDM VE interpreted as 1-RF), as does GICO over frozen BézierFlow. Native CIFAR/ImageNet use paired-image-kid-v1; GICO + BézierFlow uses paired-cifar-kid-v1 to bind the frozen two-grid transform separately. Both optimize raw uniform-minus-candidate unbiased KID differences, preserving negative estimates. Reward scaling is frozen on training evidence, balancing NFEs and classes. Complete generated/reference blocks and split isolation are required.

Only the image GICO-TF variant uses LPIPS supervision. It compares each candidate to a high-accuracy target generated by the same frozen backbone from the identical noise/class pair, using paired-lpips-v1 and raw uniform-minus-candidate VGG-LPIPS differences. This is the target-fidelity comparison with BézierFlow; the GICO + frozen BézierFlow composition above still uses KID. LPIPS and KID evidence, calibration and selection cannot be mixed or relabelled. FID50k evaluates the final generated distribution independently of either fitting objective. See image supervision for the three objective contracts.

The molecular primary reward is 40% Kabsch RMSD log improvement and 15% each for ensemble velocity, ensemble acceleration, rollout velocity and rollout acceleration norm-Wasserstein discrepancies. Metric weights apply to teacher ranking, regression, reference weighting, auxiliary scoring and reports. The fair finite-ensemble energy estimator remains a diagnostic:

mean_m ||phi(X_m) - phi(y)||
  - sum_{m != n} ||phi(X_m) - phi(X_n)|| / (2 M (M - 1))

It requires at least two independently generated complete trajectories. The frozen feature map preserves atom indices and horizon order, includes indexed pair distances and signed volumes relative to a non-collinear reference triangle, and uses deterministic atom-index tie breaking. Lengths and volumes use the training reference RMS pair distance and its cube; feature blocks use dimension normalization. The observed future is one observation, never a fabricated reference ensemble. Clashes and bond violations also remain diagnostics. The schedule report identifies energy score as a reporting diagnostic; teacher selection uses paired utility regret.

These are optimization objectives. Log transformations, mixtures of components, and teacher approximations do not establish proper scoring or benchmark improvement for the learned policy.

Shared architecture and optimization

All roles use a two-layer, width-128, four-head, pre-normalized Transformer with a 256-wide feed-forward block and task-configured dropout. Conditioning combines native context, solver identity, and continuous NFE/macro-step features. Feature normalization is fitted on training data and frozen. Initial generation noise is never a policy input.

• The teacher consumes conditioning and a candidate density and predicts the normalized metric-improvement vector. Its objective combines within-context/settings pairwise ranking (temperature 0.5) with weighted Huber regression (weight 0.25).
• The deterministic student uses 64 density-bin queries. It minimizes KL from the teacher-weighted reference-density barycenter, plus a teacher-score term.
• The stochastic student predicts 63 Gaussian log-density ratios autoregressively. Training-reference ratios are standardized; likelihood targets receive Gaussian noise with standard deviation 0.1. Predicted standard deviations are bounded to [0.05, 2]. The likelihood averages over all 63 coordinates. A reparameterized teacher-score term also trains this student.

Both students use a uniform prior over unique realized reference densities. Reference softmax temperature is expressed in paired utility units before scalar reward normalization: multiply predicted scalar scores by the frozen reward scale before dividing by temperature. Reference logits are never clipped. Stochastic target smoothing is an additional modeling choice.

Teacher-score weights are 0.01, 0.05, 0.1. Auxiliary scores use each context/solver/NFE group's frozen predicted-reference mean and population standard deviation (standard deviations below 1e-6 use 1), then clip to [-5, 5]. This normalization does not change terminal rewards or reference weights. Teacher parameters remain frozen while gradients pass through density inputs. The configurable coefficient schedule begins chasing after 60% of the student horizon; only checkpoints with a positive realized coefficient are eligible. The default remains the historical linear ramp; alternatives include a ramp with a full-weight plateau and an immediate full-weight switch.

Teacher updates average up to 64 distinct comparison groups; student updates average up to 512 distinct context/settings targets. Smaller microbatches accumulate the same equally weighted objective. Ranking pairs never cross group boundaries. Teacher checkpoint/temperature selection minimizes measured held-out reference-mixture utility regret. Context and density-family regret receive equal weight when the profile uses both; text-to-image uses context selection only. Ties prefer the configured preferred temperature, then the earlier checkpoint. Deterministic student selection maximizes frozen-teacher calibrated utility among eligible checkpoints whose held-out KL is within 15% of its eligible minimum; one density and one teacher evaluation are used per context/NFE. Stochastic selection retains measured held-out terminal utility. Both prefer earlier checkpoints on ties. Teacher checkpoints default to every 20 steps, deterministic student checkpoints to every 10, and stochastic checkpoints to every 100. See selection and score schedules.

Profile	Teacher/student steps	Dropout	Score coefficient	Initial reference temperature
Forecasting, molecules	2,000 / 2,000	0.01	0.01	0.05
SANA, SD1.5	2,000 / 2,000	0.01	0.05	0.05
CIFAR-10, ImageNet-64	2,000 / 2,000	0	0.01	0.05

Transformer AdamW defaults to learning rate 0.001 and weight decay 0.0001. Teacher and student learning rates, step limits, batch sizes and checkpoint intervals are independently configurable. Task identity comes from the evidence. All resolved settings are recorded; architecture sharing does not imply one validated temperature for every task.

The shared pool has 25 reference clocks, including late-p=3 and its reversal. Every reference is materialized through the same 64-bin representation as student outputs. Identical densities are deduplicated before mixture weighting. Historical evidence is reusable only if executed grids and measurement protocols match exactly; changed grids require new measurements.

Train and decode

The common interface accepts JSON configuration:

{
  "rows": "measurements.jsonl",
  "contexts": "contexts.npz",
  "calibration_rows": "calibration.jsonl",
  "output": "policy",
  "student_kind": "both",
  "teacher_score_weight": 0.01,
  "teacher_steps": 2000,
  "student_steps": 2000,
  "teacher_batch_groups": 64,
  "student_batch_contexts": 512,
  "microbatch_contexts": 8,
  "teacher_learning_rate": 0.001,
  "student_learning_rate": 0.001,
  "teacher_checkpoint_every": 20,
  "student_checkpoint_every": 100,
  "deterministic_checkpoint_every": 10,
  "deterministic_kl_allowance": 0.15,
  "temperatures": [0.05, 0.1, 0.5],
  "preferred_temperature": 0.05,
  "seed": 0,
  "device": "cuda",
  "purpose": "research"
}

Before fitting, add a built-in held-out utility evaluator to this configuration. Top-level paths are relative to the configuration file. Each measurement row contains task, backbone, solver, integer nfe, context_id, explicit split, integer seed, ensemble_size, reference_id, measurement_protocol, schedule_key, metrics, 64 density_mass entries, and the executed time_grid. Metric keys are crps/mase, kid for image GICO (lpips only for GICO-TF), or preference/alignment; molecular keys are kabsch_rmsd_3d, ensemble_velocity_norm_w1, ensemble_acceleration_norm_w1, rollout_velocity_norm_w1, and rollout_acceleration_norm_w1. Training input contains disjoint train and validation contexts; calibration contains only train or calibration. Locked-test rows are forbidden during fitting. Store native contexts with save_context_embedding_table.

Research molecular rows also carry the frozen molecule_feature_map dictionary from MoleculeFeatureMap.to_dict(). It is recorded in the policy artifact and checked against runtime reference geometry.

genode-train-gico --config train.json --dry-run
genode-train-gico --config train.json --student-kind both --teacher-score-weight 0.01

Research evidence requires all 25 references in every cell. Explicit purpose: functional permits a reduced reference set for integration checks; it does not produce benchmark evidence. Checkpoints are selected using validation evidence and the profile-specific teacher density-family holdout. Output directories must be new.

To compare students using an existing frozen teacher, add "teacher_artifact": "previous-policy" to the common training configuration (or pass teacher_artifact to fit). Public reuse requires selected-teacher weight/conditioning/step/temperature proof and a matching minimum-regret history entry. This skips teacher fitting and reuses its selected weights and temperature. Evidence, reward calibration, teacher conditioning and teacher fitting settings must match the source artifact; student conditioning and fitting settings can differ. The output records the source artifact checksum and inherited teacher selection history. --dry-run also checks this binding. Microbatch size can change for student memory requirements without changing the effective batch or inherited teacher.

from genode.gico.policy import load_policy

policy = load_policy("policy", student_kind="GICO-sto-policy", expected_backbone="checkpoint-sha")
grid = policy.materialize(native_context, "euler", 8, seed=412, request_id="example:member:0")

Both students use one density-to-clock implementation. It mixes in exactly 1e-8 uniform density before inverse-CDF conversion, validates representable solver grids, and enforces exact NFE accounting. The clock RNG is separate from generation noise. Sample once per generated image or trajectory; reuse each molecular member's clock throughout its rollout. Inference performs no teacher scoring, rejection, or reward-based selection.

Forecast and molecular Python evaluators accept policy and clock_seed. Forecast policy evaluation uses batch size 1, paired with the same uniform protocol. Molecular policy evaluation conditions on the initial observed history and keeps the sampled clock throughout every horizon.

Image and latent-image workflows

The separate JS-REINFORCE comparator provides a documented paper-based Dirichlet scheduler with native noise/text conditioning, budgeted SANA fitting and versioned inference artifacts. It is distinct from both GICO students. The ReFlow comparison bridge loads the original frozen EMA network and supplies a shared exact-count EDM FID evaluator.

For GICO on RF++ or EDM-as-1RF, supply paired-image-kid-v1 rows with metrics: {"kid": value}. CIFAR/ImageNet preparation validates paired feature-block evidence and native backbone/context bindings:

genode-image-gico prepare --manifest raw.json --output evidence.json
genode-image-gico train --evidence evidence.json --output policy --student-kind both --teacher-score-weight 0.01 --selection-evaluator evaluator.json
genode-image-gico validate --help
genode-image-gico materialize --help

The input schema is documented by prepare_image_rows in genode.gico.image_supervision and image supervision. The same commands fit GICO-TF only when rows explicitly declare paired-lpips-v1 with LPIPS target pairs. GICO-det-policy and GICO-sto-policy select the student architecture; the recorded image objective determines whether the run is GICO or GICO-TF. Supply global metric measurements separately when reporting them.

SANA/SD1.5 collection records executed grids, density masses, native contexts, generation seeds, and scorer identities. Use genode-latent-clock prepare-gico --help to convert paired collection evidence and independent pilot rows into the common configuration, then:

genode-latent-clock fit-gico --config evidence/train_config.json --student-kind both
genode-latent-clock prepare-collection --help

BO, PG, and LD3 remain separate comparison methods. Completed experiments remain historical results; this architecture change does not relabel or upgrade their policy states.

Artifacts and validation

Protocol genode-gico-v6 stores policy.pt plus a checksummed manifest.json. Artifacts record architecture, reward calibration, context normalization, reference densities and executed grids, split identities, solver semantics, RNG configuration, resolved fitting profiles, explicit metric weights, dropout, temperature units, normalization/selection protocols, selected steps and realized coefficients, and fitting history. Existing valid v6 student artifacts remain loadable through the explicit role codec. Public selectors are GICO-det-policy and GICO-sto-policy; both fits both. Old selector aliases are rejected. Serialized v6 role keys and fingerprint inputs remain stable. All v5 loading and older teachers lacking selection proof require their archived runtimes. Recorded clock scope must describe one complete clock per generated trajectory.

genode-report-gico-locked-test applies an artifact's frozen calibration to paired test measurements without selection. genode-evaluate-schedule-summary performs the analogous validation report. Both require new output files and matching frozen measurement protocols, native backbone bindings, and molecular feature maps. Supply policy_sha256 and student_kind for learned-policy measurements and --contexts contexts.npz for exact selected-policy clock replay. Stochastic rows require distinct clock_seed/clock_request_id pairs for each generated member and replicate. Explicit non-GICO comparisons use measurement_role: "baseline".

Evaluation rows can carry a fixed density_mass/time_grid pair or sample_clocks, a list containing one such pair per ensemble member. This supports independently sampled stochastic clocks while averaging repeated terminal measurements before constructing log improvements. Forecast evaluators export sample_clocks; molecular evaluators export the same pairs with rollout provenance in sample_clock_records. Training reference clocks remain fixed across repeats.

ImageNet reports require all 1,000 classes in each paired panel. They average classes equally and estimate uncertainty across panels, accounting for shared reference-block noise within a panel.

python -m ruff check .
python -m ruff format --check .
python -m pytest -q
python -m pip check
python -m build
git diff --check

Tests cover paired rewards, split isolation, geometric energy scoring, causal stochastic sampling, teacher-input gradients, density decoding, solver accounting, artifact integrity, and active task routing. Portable functional examples are in examples; actual assets and environment-specific operators remain external. Fixture coverage alone does not validate a pretrained generator or establish quality gains.

Controlled fitting options

All roles share a two-layer, width-128, four-head pre-normalized Transformer with feed-forward width 256, SiLU projections, density-bin RoPE, and additive MLP conditioning applied once before the block stack. RoPE positions use fixed 64-bin centers even during stochastic prefix decoding. The teacher uses raw log masses; deterministic students use bin queries and stochastic students use shifted log-ratio prefixes. Solver identity and frozen-normalized continuous log NFE/macro-step features use the same encoding for all roles.

Profiles accept independent teacher_context_mode and student_context_mode (native, the default, or global). Global mode zeroes only that role's normalized context features. With a native teacher and global student, reference targets and auxiliary scores still use each original prompt context; only the student's input is global. Groups remain separate and equally weighted. Global inference produces one deterministic density per solver/NFE, or a context-independent stochastic clock distribution, without invoking the teacher.

The task identifies the dataset; an optional backbone identity binds a fitting configuration to its measurement evidence and is always resolved in saved artifacts. Temperature candidates, preferred temperature, teacher-score weight, dropout, learning rates and horizons remain profile-specific, with explicit run overrides. For example, a SANA configuration can set teacher_context_mode: "native", student_context_mode: "global", teacher_score_weight: 0.05, temperatures: [0.05, 0.1, 0.5], and preferred_temperature: 0.05. Select on held-out selection contexts and freeze before confirmation. These architectural choices do not themselves establish better utility or generalization. Historical v4 artifacts require their originating runtime and are not upgraded.

SD1.5 variable-step integration

SD1.5 runtime configurations default to "solver": "ipndm_v", the order-2 variable-step Adams–Bashforth solver in r = sigma / alpha, with Euler startup and exactly one denoiser evaluation per update. The runtime records solver_protocol="sigma_over_alpha_ab2_v1". Integration and model-evaluation times must match; interval ratios use r, not diffusion time or log-SNR.

{
  "backbone": "sd15",
  "solver": "ipndm_v",
  "source": "assets/LD3",
  "source_revision": "ec1bf603fb19696966ca30198ed209ae6488a3e5",
  "asset_manifest": "assets/asset-manifest.json"
}

The historical fixed-coefficient solver remains explicitly available as "solver": "ipndm" for the official LD3 baseline. These are distinct solver identities, not interchangeable artifact versions. Changing solvers requires fresh paired reference measurements, anchors and reward calibration before refitting GICO. Historical policies and results retain their original runtime; do not relabel their evidence as variable-step measurements. The shared teacher, student objectives and task-specific fitting parameters are unchanged by this solver correction.

The Python fitting API has an optional checkpoint_callback(kind, step, model, statistics) observer for diagnostics. Observers must not mutate the model or RNG state. Diagnostic checkpoints do not change production eligibility: only validation-selected student states with a positive chasing coefficient are saved as deployable artifacts.

For the KID-supervised composition with a frozen transformed sampler, see GICO on frozen BézierFlow. The separate GICO-TF comparison uses paired LPIPS target fidelity.