vike
Research

Inference

The pure-Rust tree walker that predicts from a LightGBM model text with nothing native behind it — the parser's refusals, the routing rules read out of LightGBM's own source, and the two gates that prove the walker agrees with the trainer.

vike-ml splits training from inference on purpose, and crates/vike-ml/src/lib.rs's header says why. Training drives the upstream LightGBM CLI binary as a child process, once, on the box that trains (Training). Inference is compiled into every binary that runs a strategy — including a live daemon on a box that signs real orders — and shipping a C++ library there, to predict from a model that was already trained, is a cost with no benefit. So crates/vike-ml/src/infer.rs is pure Rust: it parses LightGBM's save_model text output and walks the trees, with no native dependency at all.

That split is only trustworthy if the two agree, and the header is explicit that it is gated rather than asserted: crates/vike-ml/tests/train_infer_equality.rs runs LightGBM's own task=predict over a trained model and compares it to the walker row by row, and crates/vike-ml/tests/real_model_golden.rs freezes one such model and its outputs into the repository so the comparison re-runs on every CI run, on a machine with no LightGBM binary. Model pins covers what those gates prove that the version pins cannot.

The parser: strict on purpose

crates/vike-ml/src/parse.rs is the version=v4 text-model parser and the inference path's only file read: the trainer driver and the fit cache read files of their own, but nothing that predicts does. It reads the text format rather than dump_model's JSON because the text is what output_model= writes and input_model= reads back — the CLI's own round trip, which the C++ predictor treats as canonical. The header records the bakeoff proof: a model loaded from disk, trained one more iteration and reserialised (100 → 101 trees) gave predictions at num_iteration_predict=100 byte-identical to the original's. Doubles are written with about 17 significant digits, so the file is a lossless round trip of the in-memory model.

The header's second section is titled "Strictness, on purpose", and every refusal in it exists because the alternative is a plausible-looking number from the wrong arithmetic:

  • an unknown version= is MlError::UnsupportedVersion, never a best-effort parse;
  • a non-binary objective, a multiclass model (num_class or num_tree_per_iteration other than
    1. and a linear tree (is_linear=1) are MlError::Unsupported;
  • average_output — the bare line boosting=rf writes — is refused by name, because this walker sums trees and an averaged model would be wrong by a factor of the iteration count;
  • the end of trees terminator is required, because a file cut at a tree boundary is a byte-exact prefix of a real model in which every line is valid, and accepting it loads a shorter, silently under-boosted model;
  • the header's tree_sizes count is cross-checked against the trees actually present when the line exists — the only check that catches a tree removed from the middle;
  • check_children refuses a child index that leaves the arrays or points backwards, because the walk indexes those arrays unchecked per row and would otherwise panic or loop inside a live prediction.

load_model_file is the inference path's one I/O: it reads a path the caller named, searches nowhere, reads no environment variable and has no default location — a library that guesses where its model lives is a library its caller cannot deploy.

What the model keeps

crates/vike-ml/src/model.rs is the parsed shape reduced to what inference reads. split_gain, leaf_weight, leaf_count, internal_value, internal_weight, internal_count and feature_infos are parsed past and dropped as training bookkeeping (the first of those is what Feature importance reconstructs in a separate pass). A Tree is parallel arrays: num_leaves - 1 entries per internal node and num_leaves leaf values. A child entry >= 0 indexes another internal node; a negative one encodes a leaf as its one's complement, leaf = !node. The only Objective is Binary { sigmoid }, and the sigmoid comes from the header's objective= line — not a per-tree field, and not always 1.0.

shrinkage is dropped too, and the module calls that one a trap rather than a tidy-up.

The routing rules

infer.rs's header states that every routing rule was read out of LightGBM's include/LightGBM/tree.h rather than inferred, because each one fails silently: the walk terminates whatever you do, at a different leaf, and the probability that comes out looks like a probability.

numerical_decision is three rules in order. A NaN at a node whose missing_type is not Nan is remapped to 0.0 and compared normally; a value that is the node's missing value goes to the default_left side; otherwise fval <= threshold goes left — <=, not <. missing_type is the two bits above the low two of decision_type, and MissingType::Zero uses LightGBM's kZeroThreshold, pinned in model.rs as ZERO_THRESHOLD = 1e-35f32 as f64 because upstream's literal is a float widened to double.

categorical_decision is a different function, and its doc calls it the place a hand-written walker is most likely to diverge in production. It never consults default_left or missing_type: a NaN or a negative category goes right unconditionally. The category id is fval as i32 — truncation toward zero, not rounding — the node's threshold is an index into cat_boundaries, not a value, and membership is a bit test in a packed 32-categories-per-word bitset where a word past the group's end is implicitly zero.

leaf_value_for returns leaf_value[!node]; a constant tree (empty split arrays) returns its only leaf without walking.

The score, and the three things it refuses to do

GbdtModel::raw_score is the plain sum of every tree's selected leaf value, and its doc lists what it deliberately does not do. It does not multiply by shrinkage: LightGBM scales leaf_value by the learning rate in place at training time, so applying it again scales every score by the learning rate squared. It does not add a base score: boost_from_average's init score is folded into the first tree (AddBias), so summing every tree in file order is both necessary and sufficient. And it does not use a compensated sum — vike_model::py_sum is more accurate, which is exactly why it is wrong here: the equality gate's the_raw_scores_match_bit_for_bit compares this function against LightGBM's own predict_raw_score=true output bit for bit.

predict_proba is 1 / (1 + exp(-sigmoid * raw)), and its exp is libm::exp, not f64::exp. The doc calls that the one place a probability out of this crate could differ between two boxes: IEEE 754 requires nothing of exp, so a method call reaches whichever libm the platform ships, and the two gates over this function each budget 1e-12 for one exp through two libms.

The single-row API is unchecked; the batch API is not

raw_score and predict_proba validate nothing in release — they are the per-tick live path — and Tree::feature_value substitutes NaN for an entry the row does not carry, which diverges from LightGBM's own 0.0 fill at a missing_type == Nan node. A debug_assert! on the row width catches drift in test and debug builds; predict_proba_batch (and its pre-transform twin raw_scores_batch) is the checked entry point, refusing with MlError::Shape a width the model cannot use or a buffer that is not a whole number of rows.

Where the probabilities go once a study has them is Models; the tutorial that fits one end to end is Train a model.

On this page