FlowStateModel

FlowStateModel at 6 layers and hidden size 512. One published checkpoint has this shape.

Implementing it

How to read this sheet: method.

No forward was recorded for this shape, so the block is read from the published configuration, the checkpoint's tensor names, and modeling_flowstate.py in granite-tsfm v0.3.9 — not transcribed from a run. Keys the configuration omits take the code defaults: mlp_type: "gated", out_gating: false, scale_factor: 1.0.

Generated from granite-timeseries-flowstate-r1.

# Continuous time, and the only model here whose state is complex.
# No patching anywhere in this file.

# 1. The normalisation is causal and it scales the value only.
#    Running statistics per position: mu[k] is the mean of 0..k, and the
#    deviation's k-th term uses the mean as of its own position.
v = last 2048 samples
flag = isnan(v); v = nan_to_num(v, 0)
mu, sd = running_mean_std(v, over=~flag)  # per position, not per window;
                                          #   variance clamped at 1e-5
z = (v - mu) / sd * ~flag                 # the flag itself is not scaled
h = stack(z, flag) @ embed.embed.T        # Linear 2 -> 512

# 2. 6 S5 layers, each a scan and then a gated activation.
s = scale_factor                          # 1.0 unless the caller sets it
for layer in encoder.layers:
    lam  = -exp(ssm.log_Lambda_real) + 1j * ssm.Lambda_imag
    Lbar = exp(lam * exp(ssm.log_Delta) * s)          # discretised per step
    Bbar = (Lbar - 1) / lam * (ssm.B_tilde_r + 1j * ssm.B_tilde_i)
    x = 0
    for k in 0 .. T-1:
        x = Lbar * x + Bbar @ h[k]                    # the reference computes
                                                      # this as an FFT
                                                      # convolution
    y = ssm.C_tilde_r @ Re(x) + ssm.C_tilde_i @ (-Im(x)) + ssm.D * h
    #     ^ note the sign: the real part of a complex product written out.
    y = selu(y)
    y = y * sigmoid(y @ layer.out.T)                  # self-gate
    h = h + layernorm(y, layer.norm)                  # norm the branch,
                                                      #   then add

# 3. Only the last position reaches the decoder: the last layer keeps
#    only the final state, the ones before it produce the whole sequence
#    because the next layer reads all of it.
c = h[-1] @ decoder.lin.T                 # -> 9 x 256 Legendre coefficients,
                                          #    one curve per quantile 0.1 .. 0.9
t = -1 + (1 .. 24) * 1.95 * s / 24        # `decoder_patch_len: 24`, range
                                          #   [-1, 0.95]
q = (legendre_basis(t, degree=255) / 4) @ c
q = q * sd[-1] + mu[-1]                   # denormalised, with the last
                                          #   position's statistics
# FlowStateModel returns these nine curves.

FlowStateForPrediction, the wrapper around this class, then turns the nine curves into outputs: torch.quantile over the nine at each horizon point, at the same nine probabilities (sort and linear interpolation — the identity while the curves do not cross), returned as quantile_outputs; and a point forecast, returned as prediction_outputs, that is a triangular weighted sum of those, w = 0.5 - |0.5 - p| normalised — not the median row and not a uniform mean.

Sources, granite-tsfm v0.3.9: FlowStateCausalRevIN; the missing flag and embedding with Linear(2, 512); the S5 discretisation, FFT convolution and readout; the layer; the decoder called with one 24-point patch; the quantile and point step in FlowStateForPrediction.forward.

The published class returns nine curves; the point forecast is the wrapper's. decoder.lin is 512 -> 256 * 9 and there is no tenth set of coefficients. The wrapper resamples the nine and takes a triangular weighted sum; a uniform average of the nine, or the 0.5 row, is a different number.

The denormalisation happens before the resampling, so the curves are ordered in the output's units rather than the model's. With a single positive scale the two orders agree.

The normalisation is causal and it scales the value only. Position k is scaled by the statistics of 0..k, the missing flag passes through unscaled, and the forecast is denormalised with the last position's statistics. A window-wide mean — what the other forecasters here use — runs and is a different model.

The residual adds a normalised branch. Each layer computes h + LayerNorm(gate(selu(ssm(h)))); normalising after the add is the plausible misreading.

The input is [value, missing] pairs, so the embedding is Linear(2, 512); a port that feeds bare values has the wrong input width.

The diagonal scan and the reference's FFT convolution compute the same thing: for a diagonal system the flipped kernel is Lbar^j, the FFT is zero-padded to twice the length so the convolution is linear rather than circular, and a causal convolution with that kernel is the closed form of the recurrence. They agree up to floating-point rounding.

Geometry

How this was checked

Status and evidence grade in the assessment (see what the grades mean):

• granite-timeseries-flowstate-r1: runs, class — the architecture class is implemented; nothing checkpoint-specific was measured

The implementation meant here and below is an unpublished independent inference implementation by the author.

What stands behind the block above, beyond the published configuration it is read from:

• a checkpoint loader, which is where published tensor names are read
• an independent reference implementation of this architecture in Python, driving the published modelling code — an executable statement of what the model should compute, written against the publication rather than against any one implementation of it

Where an independent implementation and the published configuration disagree, the configuration is what this page reports and the disagreement is what it says.

Checkpoints with this architecture