Data-Driven Invariant Learning for Probabilistic Programs (Extended Abstract)

1 Introduction

Probabilistic programs are standard imperative programs augmented with a sampling command-a mechanism to draw random samples from a probability distribution. Probabilistic programs provide a formal way to describe randomized computations. While the mathematical semantics of such programs is fairly well-understood [Kozen, 1981], verification methods remain an active area of research. Existing automated techniques are either limited to specific properties (e.g., [Smith et al., 2019; Albarghouthi and Hsu, 2018; Carbin et al., 2013; Roy et al., 2021]), or target simpler computational models [Baier et al., 1997; Kwiatkowska et al., 2011; Dehnert et al., 2017].

Reasoning about Expectations. One of the earliest methods for reasoning about probabilistic programs is through expectations. Originally proposed by Kozen [Kozen, 1985], expectations generalize standard, binary assertions to quantitative, real-valued functions on program states. Morgan and McIver further developed this idea into a powerful framework for reasoning about probabilistic imperative programs, called the weakest pre-expectation calculus [Morgan et al., 1996; McIver and Morgan, 2005]. The weakest pre-expectation calculus defines the weakest pre-expectation (wpe) operator that takes an expectation E and a program P to produce an expectation E′ such that E′ (σ) is the expected value of E in the output distribution JPKσ. In this way, the wpe operator can be viewed as a generalization of Dijkstra's weakest pre-conditions calculus [Dijkstra, 1975] to probabilistic programs. The wpe operator has two key strengths: first, it enables reasoning about probabilities and expected values; second, when P is a loop-free program, it is possible to transform wpe(P,E) into a form that does not mention the program P via simple, mechanical manipulations, essentially analyzing the effect of the program on the expectation through syntactically transforming E. However, there is a caveat: the wpe of a loop is defined as a least fixed point, and it is generally difficult to simplify this into a more tractable form. Fortunately, the wpe operator satisfies a loop rule that simplifies reasoning about loops: if we can find an expectation I satisfying an invariant condition, then we can easily bound the wpe of a loop. Checking the invariant condition involves analyzing just the body of the loop, rather than the entire loop. Thus, finding invariants becomes the primary bottleneck towards automated reasoning about probabilistic programs.

Our Approach. Our approach to synthesizing loop invariants for probabilistic programs is inspired by data-driven invariant learning techniques [Flanagan and Leino, 2001; Ernst et al., 2007] for regular programs. In these methods, the program is executed with a variety of inputs to produce a set of execution traces. This data is viewed as a training set, and a machine learning algorithm is used to find a classifier describing the invariant. Data-driven techniques reduce the reliance on templates, and can treat the program as a black box-the precise implementation of the program need not be known, as long as the learner can execute the program to gather input and output data. But to extend the data-driven method to the probabilistic setting, there are a few key challenges:

• Quantitative invariants. While the logic of expectations resembles the logic of standard assertions, an important difference is that expectations are quantitative: they map program states to real numbers, not a binary

yes/no. While invariant learning for deterministic programs is a classification task (i.e., predicting a binary label given a program state), our probabilistic invariant learning is closer to a regression task (i.e., predicting a number given a program state).

• Stochastic data. Invariant learning for deterministic programs assumes that the program behaves like a function: a given input state always leads to the same output state. In contrast, a probabilistic program takes an input state to a distribution over outputs. Since we are only able to observe a single draw from the output distribution each time we run the program, execution traces in our setting are inherently noisy. Accordingly, we cannot hope to learn an invariant that fits the observed data perfectly, even if the program has an invariant-our learner must be robust to noisy training data.
• Complex learning objective. To fit a probabilistic invariant to data, the logical constraints defining an invariant must be converted into a regression problem with a loss function suitable for standard machine learning algorithms and models. While typical regression problems relate the unknown quantity to be learned to known data, the conditions defining invariants are somehow self-referential: they describe how an unknown invariant must be related to itself. This feature makes casting invariant learning as machine learning a difficult task.

Contributions. The contributions of our work are:

• We provide a general algorithm, EXIST (EXpectation Invariant SynThesis), for learning invariants for probabilistic programs. EXIST executes the programmultiple times on a set of input states, and then uses machine learning algorithms to learn models encoding possible invariants. A counterexample guided inductive synthesis (CEGIS) loop iteratively expands the dataset after encountering incorrect candidate invariants.
• We describe instantiations of EXIST tailored for handling two problems: learning exact invariants, and learning sub-invariants. Our method for exact invariants learns a model tree [Quinlan, 1992], a generalization of binary decision trees to regression. The constraints for sub-invariants are more difficult to encode as a regression problem, and our method learns a neural model tree [Yang et al., 2018] with a custom loss function. While the models differ, both algorithms leverage offthe-shelf learning algorithms.
• We evaluate our implementation of EXIST on a large set of benchmarks. EXIST learns invariants for examples considered in prior work and new, more difficult versions that are beyond the reach of prior work.

This article is the extended abstract of the conference paper [Bao et al., 2022a] published in CAV'22. An extended version of the article is also available [Bao et al., 2022b].

2 Preliminaries

Probabilistic Programs. We will consider programs written in pWhile, a basic probabilistic imperative language with the following grammar, where e is a boolean or numerical expression.

All commands P are interpreted into maps from memories to distributions over memories [Kozen, 1981]. We write JP Kσ for the output distribution of program P from initial state σ. Since we will be interested in running programs on concrete inputs, we will assume throughout that all loops are almost surely terminating; this property can often be established by other methods (e.g., [Chatterjee et al., 2016a; Chatterjee et al., 2016b; McIver et al., 2018]).

Weakest Pre-expectation Calculus. Morgan andMcIver's weakest pre-expectation transformer (wpe) takes a program P and an expectation E to another expectation E′ , called the pre-expectation; wpe is defined in Fig. 1. The case for loops involves the least fixed-point (lfp) of Φwpe E := λX.([e] · wpe(P,X)+[¬e] ·E), the characteristic function of the loop with respect to wpe [Kaminski et al., 2016]. The characteristic function is monotone on the complete lattice E, so the least fixed-point exists by the Kleene fixed-point theorem. The key property of the wpe transformer is that for any program P, wpe(P, E)(σ) is the expected value of E over the output distribution JPKσ. Intuitively, the weakest pre-expectation calculus provides a syntactic way to compute the expected value of an expression E after running a program P, except when the program is a loop. For a loop, the least fixed point definition of wpe(while e : P, E) is hard to compute.

3 Problem Statement

Analogous to when analyzing the weakest pre-conditions of a loop, knowing a loop invariant or sub-invariant expectation enables one to easily bound the loop's weakest preexpectations. However, (sub)invariant expectations can be difficult to find. We develop an algorithm to synthesize invariants and sub-invariants of probabilistic loops. More specifically, our algorithm tackles the following two problems:

1. Finding exact invariants: Given a loop while G : P and an expectation postE as input, we want to find an expectation I such that

Such an expectation I is an exact invariant of the loop with respect to postE. Since wpe(while G : P, postE) is a fixed point of Φwpe postE , wpe(while G : P, postE) has to be an exact invariant of the loop. Furthermore, when while G : P is almost surely terminating and postE is upper-bounded, the existence of an exact invariant I implies I = wpe(while e : P, E).

1. Finding sub-invariants: Given a loop while G : P and expectations preE, postE, we aim to learn an expectation I such that

Figure 1: Morgan and McIver's weakest pre-expectation transformer (wpe)

The first inequality says that I is a sub-invariant: on states that satisfy G, the value of I lower bounds the expected value of itself after running one loop iteration from initial state, and on states that violate G, the value of I lower bounds the value of postE. Any sub-invariant lower-bounds the weakest pre-expectation of the loop, i.e., I ≤ wpe(while G : P,E) [Kaminski, 2019]. Together with the second inequality preE ≤ I, the existence of a sub-invariant I ensures that preE lowerbounds the weakest pre-expectation.

4 Methodology

Our data-driven method runs a Counterexample Guided Inductive Synthesis (CEGIS), but differs from conventional CEGIS tools in two ways. First, candidates are synthesized by fitting a machine learning model to data consisting of program traces starting from random input states. Our target programs are also probabilistic, introducing a second source of randomness to program traces. Second, our approach seeks high-quality counterexamples-violating the target constraints as much as possible-in order to improve synthesis. For synthesizing invariants and sub-invariants, such counterexamples can be generated by using a computer algebra system to solve an optimization problem.

EXIST takes a probabilistic program P, a post-expectation or a pair of pre/post-expectation pexp, to produce a loop (sub)invariant expectations. The user can provide two hyperparameters, Nruns and Nstates , to control the data-generation process. With these inputs, EXIST proceeds below:

Generate Features. EXIST starts by generating a list of features feat, which are numerical expressions formed by program variables used in P.

Sample Initial States. Next, EXIST samples Nstates number of initial states by uniformly sampling values of program variables from their respective domains.

Sample Training Data. For learning exact invariants, from each initial state si, EXIST runs P until termination for Nruns times to get the list of final states {σ1, . . . , σNruns } and then produces the training example:

Above, the value vi is the empirical mean of postE in the output state of running P from initial state si; as Nruns grows large, this average value approaches the true expected value wpe(P, postE)(si).

For learning subinvariants, from each si, EXIST runs single iteration of P (and then restart at si) for Nruns times to get a list of output states {σ1, . . . , σNruns } and then produces the training example:

(Machine) Learn an Invariant. In each iteration of the CEGIS loop, first the learner learnInv trains models to minimize violation of the required inequalities (i.e., Eq. (1) for learning exact invariants; Eqs. (2) and (3) for learning subinvariants) on data.

For learning exact invariants, we choose regression models to be model trees, which are decision trees with customizable models instead of labels on leaves. While our methods can potentially work for other leaf models, we focus on linear models or multiplicative models (which are linear models on the logarithm space of the data) because of their simplicity and expressiveness. This class of model trees suits our goal because they can be easily translated into numerical expressions, which are usually the form people use to encode expectations. With the training data = {(s1, v1), . . . , (sNstates , vNstates )}, where each vi approximates wpe(P, postE)(si), we train a model tree T that takes the feature vector of si, denoted F(si), as input and predicts vi. We use the standard mean-square-error to measure the error between predicted values T(F(si)) and the target value vi and apply off-the-shelf tools for training.

For learning sub-invariants, we initially also want to use model trees, but the loss is more complicated and the standard model tree training mechanism based on divide-and-conquer no longer applies. A remedy is to apply gradient descent to train on that loss, but standard model trees are not differentiable. To bridge the gap, EXIST trains neural model trees [Yang et al., 2018], which are neural networks for approximating model trees, to fit the data and then translate the learned neural networks back to model trees.

Extract Expectation Invariants from Models. EXIST now translates the learned models into numerical expressions and regards them as the set of candidate(sub)invariants.

Verification. For each generated candidate invariant inv, EXIST attempts to verify if inv satisfies the required constraints. If it cannot find any program state where inv violates the required set of constraints, the verifier returns inv as a valid invariant (or sub-invariant).

Data Augmentation. If inv violates any inequalities, EXIST produces a set of counterexample program states that are added to the set of initial states. To ensure that the generated counterexamples are effective data-points, EXIST looks for program states that maximize inv's violation of required inequalities. EXIST then adds states in counterexamples to states and augments the existing data with new data points generated on these counterexample states. The data augmentation process ensures that the synthesis algorithm collects more and more initial states, some randomly generated (sampleStates) and some from prior counterexamples, guiding the learner towards better candidates.

EXIST repeats the CEGIS loop by re-learning new models on the augmented dataset. Like most of the other CEGISbased tools, our method is sound but not complete, i.e., if the algorithm returns an expectation then it is guaranteed to be an exact invariant or sub-invariant, but the algorithm might never return an answer; in practice, we set a timeout to force the procedure to terminate.

5 Evaluations

We implemented our prototype in Python, using sklearn and tensorflow to fit model trees and neural model trees, and Wolfram Alpha to verify and perform counterexample generation. We have evaluated our tool on a set of 18 benchmarks drawn from different sources in prior work [Gretz et al., 2013; Chen et al., 2015; Kaminski and Katoen, 2017]. We summarize our findings as follows:

• EXIST successfully synthesized and verified exact invariants for 14/18 benchmarks within a timeout of 300 seconds. Our tool was able to generate these 14 invariants in reasonable time, taking between 1 to 237 seconds. The sampling phase dominates the time in most cases. We also compare EXIST with a tool from prior literature, MORA [Bartocci et al., 2020]. We found that MORA can only handle a restrictive set of programs and cannot handle many of our benchmarks. We also discuss how our work compares with a few others in (Section 6).
• To evaluate sub-invariant learning, we created multiple problem instances for each benchmark by supplying different pre-expectations. On a total of 34 such problem instances, EXIST was able to infer correct invariants in 27 cases, taking between 7 to 102 seconds.

Table 1 and Table 2 contain some of the exact invariants and sub invariants generated by EXIST.

6 RelatedWork

Invariant Generation for Probabilistic Programs. The PRINSYS system [Gretz et al., 2013] employs a templatebased approach to guide the search for probabilistic invariants. [Chen et al., 2015] proposed a counterexample-guided approach to find polynomial invariants, by applying Lagrange interpolation. Unlike PRINSYS, this approach does not need templates; however, invariants involving guard expressionscommon in our examples-cannot be found, since they are not polynomials. Additionally, [Chen et al., 2015] uses a weaker notion of invariant, which only needs to be correct on certain initial states; our tool generates invariants that are correct on all initial states. [Feng et al., 2017] improves on [Chen et al., 2015] by using Stengle's Positivstellensatz to encode invariants constraints as a semidefinite programming problem. However, their approach cannot synthesize piecewise linear invariants.

Another line of work applies martingales to derive insights of probabilistic programs. [Chakarov and Sankaranarayanan, 2013] showed several applications of martingales in program analysis, and [Barthe et al., 2016] gave a procedure to generate candidate martingales for a probabilistic program; however, this tool gives no control over which expected value is analyzed - the user can only guess initial expressions and the tool generates valid bounds, which may not be interesting.

Data-driven Invariant Synthesis. We are not aware of other data-driven methods for learning probabilistic invariants, but a recent work [Abate et al., 2021] proves probabilistic termination by learning ranking supermartingales encoded as two-layer neural networks from trace data.

Data-driven inference of invariants for deterministic programs has been an area of interest, starting from DAIKON [Ernst et al., 2007]. For inductive invariants, ICE learning [Garg et al., 2016] uses a modified decision tree learning algorithm and HANOI [Miltner et al., 2020] uses a CEGIS-based engine that alternates between weakening and strengthening candidates. Recent work uses neural networks to learn invariants [Si et al., 2018]. Data from fuzzing has been used for almost correct inductive invariants [Lahiri and Roy, 2022] for programs with closed-box operations. As a data-driven learning task, invariant inference for deterministic programs can be treated as a classification problem while that for probabilistic programs becomes a regression task.