O PTIMAL D ATA S AMPLING FOR T RAINING N EURAL S URROGATES OF P ROGRAMS
Anonymous authors
Paper under double-blind review
A BSTRACT
Programmers and researchers are increasingly developing surrogates of programs, models of a subset of the observable behavior of a given program, to solve a variety of software development challenges.
Programmers train surrogates from measurements of the behavior of a program on a dataset of input examples.
We present a methodology for optimally sampling datasets to train neural network based surrogates of programs.
We ﬁrst characterize the optimal proportion of data to sample from each path in a program based on the complexity of learning the path.
We next provide a program analysis to determine the complexity of different paths in a program.
We evaluate these results on a large-scale graphics program, demonstrating that theoretically optimal sampling results in empirical improvements in accuracy.
I NTRODUCTION
Programmers and researchers are increasingly developing surrogates of programs, models of a subset of the observable behavior of a given program, to solve a variety of software development challenges (Renda et al., 2021).
For example, Esmaeilzadeh et al.
(2012) train small neural networks to mimic existing programs, then deploy the neural networks in place of the programs to speed up computation.
Generally, surrogates are used to accelerate programs (Esmaeilzadeh et al., 2012; Mendis et al., 2019; Munk et al., 2019), apply transfer learning to programs (Tercan et al., 2018; Kustowski et al., 2020; Kwon & Carloni, 2020), and approximate the gradient of programs to optimize their inputs (Renda et al., 2020;
She et al., 2019; Tseng et al., 2019).
Dataset generation.
Training a surrogate of a program requires measurements of the behavior of the program on a dataset of input examples.
There are three common approaches to collecting this dataset.
The ﬁrst is to use data that is uniformly sampled (or sampled using another manually deﬁned distribution) from the input space of the program (Tseng et al., 2019; Kustowski et al., 2020).
The second is to use data instrumented from running the original program on a workload of interest (Renda et al., 2020; Esmaeilzadeh et al., 2012).
The third is to use active learning (Settles, 2009), a class of online methods that iteratively query labels for data points based on the expected improvement in accuracy resulting from additional samples (˙Ipek et al., 2006;
She et al., 2019; Pestourie et al., 2020).
These approaches show promise, but they face challenges with programs with control ﬂow.
Programs with control ﬂow (e.g. branches and loops) are piecewise functions: each control ﬂow path induces a different trace of operations that are applied to the input.
The sampling techniques above do not optimally allocate samples between different paths, resulting in surrogates which do not adequately learn the behavior of the program along all paths.
For example, Renda et al.
(2020, Section IV.A) identify a scenario in which an instrumented dataset does not exercise a set of paths in the program enough times for the surrogate to learn the behavior along those paths.
Our approach.
Our approach uses the source code and semantics of the program under study to guide dataset generation for training a surrogate of the program.
The core concept is to analyze the complexity of each path in a program and to allocate more samples to paths that are more complex to learn.
Stratiﬁed functions.
Our approach represents the program as a stratiﬁed function , a function with different behavior in different regions ( strata ) of the input space (i.e., a piecewise function).
1 We use stratiﬁed surrogates to model such functions.
To construct a stratiﬁed surrogate, we train independent surrogates of each component of the stratiﬁed function.
At evaluation time, a stratiﬁed surrogate checks which stratum an input is in (using the original program) then applies the corresponding surrogate.
This evaluation-time stratum check must not preclude the use of the surrogate for its downstream task.
We therefore adopt a standard modeling assumption in the approximate computing literature: that precisely determining paths is an acceptable cost during approximate program execution
(Sampson et al., 2011; Carbin et al., 2013).
Optimal sampling.
With this stratiﬁed modeling assumption, we then determine how many samples to allocate to train each surrogate.
Using neural network sample complexity bounds for learning analytic functions (Arora et al., 2019; Agarwala et al., 2021)
we calculate a complexity for each component functionwhichgivesanupperboundonhowmanysamplesarerequiredtolearnthebehaviorofthatcomponenttoagivenerror.
Givenadatadistributiondescribingthefrequencyofeachcomponentandgiven eachcomponentfunction’scomplexity,wethenderivetheoptimalnumberofsamplestoallocatetotrainingeachsurrogateofeachcomponent,minimizingtheupperboundonthestratiﬁedsurrogate’serror.
Complexity analysis.
We present a programming language, T URACO , in which programs denote stratiﬁed functions with well-deﬁned complexity measures.
We provide a program analysis for
T URACO programs that automatically determines the strata of the function and calculates an upper bound on the complexity of each component of the stratiﬁed function that the program denotes.
Renderer demonstration.
To demonstrate that optimal sampling using our complexity analysis improves surrogate accuracy on downstream tasks, we present a case study of learning a surrogate of a renderer in a video game engine.
We show that our optimal sampling approach results in between 15% and 47% lower error than training using distributions that do not take into account path complexity.
These accuracy improvements correlate with perceptual improvements in the generated renders.
Contributions.
In sum, we present the following contributions:
• An optimal approach to allocating samples among strata to train stratiﬁed neural network surrogates of stratiﬁed analytic functions that minimizes the upper bound on the surrogate’s error.
• A programming language, T URACO , in which all programs are learnable stratiﬁed functions, and a program analysis to determine the complexity of learning surrogates of those programs.
• An evaluation of these results on a graphics program, demonstrating that theoretically optimal sampling using T URACO ’s complexity analysis results in empirical improvements in accuracy.
We lay the groundwork for analyzing optimal sampling approaches for training surrogates of programs.
Our results hold out the promise of surrogate training approaches that intelligently use the program’s semantics to guide the design and training of surrogates of programs.
E XAMPLE
Figure 1a presents an example distilled from our evaluation (Section 5) that we use to demonstrate how optimal path sampling , sampling from paths according both to their frequency in a data distribution and to their complexity, results in a more accurate surrogate than frequency-based path sampling , sampling according to the frequency of paths alone.
Program under study.
We study a graphics program that calculates the luminance (i.e., brightness) at a point in a scene as a function of sunPosition , the height of the sun in the sky (i.e., the time of day) which ranges from − 1 to 1 , and emission , a property of the material at that point which ranges from − 1 to 1 .
The program ﬁrst checks whether it is daytime (Line 2), and sets the ambient lighting variable accordingly.
The program next checks whether the sun position is above a threshold (Line 7) and sets the emission variable accordingly.
The output is then the sum of the ambient light and the light emitted by the material.
Figure 1b presents the output of this program on inputs between − 1 and 1 . [Figure] luminance
[Figure]
properties.
0 sunPosition - 1 .- 0 .- 0 .0 .0 .
.emission
Luminance Across Input Distribution 0 .0 .1 .1 .2 .
L u m i n a n ce
[Figure] [Figure]
: sunPosition
[Figure]
fun ( sunPosition , emission ) {
[Figure]
=
[Figure]
emission * =
[Figure]
[Figure]
[Figure]
ambient + emission [Figure] [Figure] sunPosition
[Figure]
sunPosition [Figure] sunPosition [Figure]
sunPosition
[Figure]
sunPosition [Figure] The path conditions (Lines 2 and 7) partition the program into three traces: nighttime, when sunPosition is less than 0 (Figure 1c); twilight, when sunPosition is between 0 and 0 .(Figure 1d); and daytime, when sunPosition is greater than 0 .
1 (Figure 1e).
These paths are separated by dashed black lines in Figure 1b.
Training a surrogate of this program poses a particular challenge because of the different behavior of these traces.
Furthermore, these traces have different relative complexities: when sunPosition is less than 0 .
1 the function is linear, but when sunPosition is above 0 .
1 the function is quadratic.
We must ensure that the data distribution that we use for training surrogates reﬂects not only the different paths of the program, but also the relative complexities of each path of the program.
Optimal path sampling.
We present an approach to determining the optimal amount of training data to sample from each path to train a stratiﬁed surrogate of this program.
Speciﬁcally, given a data distribution and a data budget we want to ﬁnd the optimal number of data points to sample from each path to minimize the expected error of a surrogate of the program over the data distribution.
Intuitively, our approach is to prioritize sampling paths that are frequent and paths that are complex (and thus require more samples to learn).
First we determine the frequency of each path in the underlying data distribution.
We assume that the data has a uniform distribution over inputs between − 1 and 1 .
This results in path frequencies for the nighttime path ( sunPosition < 0 ) of 50% , the twilight path ( 0 < sunPosition < 0.1 ) of 10% , and the daytime path ( 0.1 < sunPosition ) of 40% .
Next we determine the sample complexity of each path, the number of samples required to learn the function along that path to a given error.
We use the sample complexity results of Agarwala et al.
(2021), who give an upper bound on the number of samples required to learn a neural network approximation of a given function.
Using this bound (as implemented by our T URACO analysis described in Section 4.2), we determine that the twilight path takes 1 .
5 × as many samples to train a surrogate to a given error as the nighttime path, and the daytime path requires 5 × as many samples.
Finally, we combine these metrics to determine the optimal sampling rates.
Using Equation (3) in
Section 3.2, we ﬁnd that the optimal sampling rate is to sample 39 .
3% of the data from the nighttime path, 14 .
3% of the data from the twilight path, and 46 .
4% of the data from the daytime path.
Stratiﬁed surrogates.
The class of surrogate model for which the above approach is optimal is that of a stratiﬁed neural surrogate – a set of disjoint neural networks which are applied based on which
[Figure]
path the inputs induce in the program.
Concretely, this means that we train one surrogate per path, and pick which to apply for each input at evaluation time.
For this example program, picking which surrogate to apply just requires comparing sunPosition against constant threshold values.
Training methodology.
For each surrogate, we train a 1-hidden-layer MLP with 512 hidden units with a ReLU activation, using 10,000 steps of Adam with learning rate 0 .
0005 and batch size 128.
Results.
Figure 2 presents the error of surrogates trained to mimic this program, as a function of the training dataset size.
On the left, Figure 2a presents the error of surrogates of each path.
On the right,
Figure 2b presents the error of stratiﬁed surrogates of the entire program for optimal path sampling and a baseline of sampling according to path frequency alone.
The x axis of each plot is the dataset size used to train the surrogate (on the left, the dataset size per path; on the right, the total dataset size used for all paths).
The y axis of each plot is the error of the resulting surrogate (lower is better).
Figure 2a shows that though the complexity measure is not exactly proportional to the empirical sample complexity, it does correlate with it: at a given data budget, the daytime path with the highest complexity has the highest error, followed by twilight then nighttime.
Figure 2b shows that the optimal path sampling approach results in lower error than sampling according to path frequency alone.
For datasets of total size below 70 samples, the surrogate trained with optimal path sampling has a geometric mean decrease in error of 27 .
For datasets of total size abovesamples, the surrogate trained with optimal path sampling has a geometric mean decrease in error of 5 . 5% .
Across the entire range of dataset sizes evaluated in this plot, the surrogate trained with optimal path sampling has a geometric mean decrease in error of 15% .
O PTIMAL S AMPLING
In this section, we formally deﬁne stratiﬁed functions and stratiﬁed surrogates, and derive the optimal sampling distribution to use when training a stratiﬁed surrogate of stratiﬁed function.
3. S ETUP
We deﬁne a learning algorithm , a function that trains a surrogate of a given input function, as a random function tr :( X →Y ) ×D× N ×L→ ( X →Y ) that takes a function f : X →Y from inputs x ∈X to outputs y ∈ Y , a distribution D ∈ D over inputs x , a number of training examples n ∈ N , and a loss function (cid:96) ∈L : Y×Y
→ R ≥ 0 which measures the cost of an incorrect prediction, and returns a function (representing the output surrogate)
ˆ f : X →Y .
Equation 1 deﬁnes a given function f as probably approximately correctly learnable (abbreviated as learnable for the remainder of the paper) for a given learning algorithm tr and loss function (cid:96) if for all distributions D , with high probability 1 − δ the learning algorithm returns a surrogate ˆ f that approximately
[Equation]
Following Arora et al.
(2019) and Agarwala et al.
(2021) we study functions and learning algorithms for which there is a measure of the complexity ζ ( f ) of function f for the learning algorithm tr such that the relationship between n , ζ ( f ) , (cid:15) , and δ in Equation (1) is:
[Equation]
We instantiate the complexity measure ζ ( f ) for neural networks in Section 3.3.
We deﬁne a stratiﬁed function f as follows:
[Equation]
where c is the number of strata, { s i } ci =1 are strata, ∀ i (cid:54)
=
j.s i ∩ s j = ∅ , and ∪ i s i
= X .
We deﬁne a stratiﬁed surrogate ˆ f as a stratiﬁed function with components ˆ f i .
For a data distribution D , let D ( x ) be the probability that x is sampled from D , and D ( s i ) be the total probability mass of all data points within s i over D (i.e., D ( s i ) =
(cid:82) x ∈ s
i D ( x ) ).
Let D ( x | s i ) be the probability of a data point x sampled from D if x ∈ s i .
O PTIMAL S AMPLING
Our goal is to learn a stratiﬁed surrogate
ˆ f of a stratiﬁed function f , where each component function f
i is learnable.
We are given a data distribution D , a maximum sample budget n , a learning algorithm tr , a loss function (cid:96) , and a failure probability δ .
Our task is therefore to ﬁnd the number of samples to allocate to each stratum to train a surrogate of that function component.
We assume that each surrogate component’s failure probability δ i = δc , which satisﬁes the overall failure probability by union bound.
Note that Equation (2) is an upper bound and not an exact equality.
We are therefore minimizing the upper bound of the error of the resulting surrogate, rather than directly minimizing the error.
Formally, we solve the following optimization problem:
[Equation]
Theorem 3.1.
With δ i = δc , the upper bound of Equation (3) is minimized at:
[Equation]
The proof of this theorem is presented in Appendix C.
3. N EURAL N ETWORK L EARNABILITY
Equation (2) deﬁnes the required sample complexity for learning ˆ f as a function of ζ ( f ) , the complexity of f .
This section deﬁnes ζ ( f ) for training neural network surrogates of analytic functions.
This section is an abridged summary of assumptions and results presented by Agarwala
Agarwala et al.
(2021) provide a calculus for learning surrogates of analytic functions f (around 0) based on the tilde ˜ f of the function:
[Equation]
Note the following properties for x ≥ 0 :
[Equation]
For a 2-layer neural network trained with stochastic gradient descent, if f is analytic, (cid:126)x is on the d -dimensional unit sphere ( (cid:126)x ∈ S d ), β ∈ R d (a parameter set to control the scale of the inputs), and (cid:96) is 1-Lipschitz, then f ( β · (cid:126)x ) is learnable with:
[Equation]
T URACO : P ROGRAMS AS S TRATIFIED F UNCTIONS
In this section we present T URACO , a programming language in which all programs denote learnable stratiﬁed functions.
We provide a program analysis for T URACO programs which calculates an upper bound on the complexity of each component of the stratiﬁed functions that the program denotes.
4. S YNTAX AND S TANDARD I NTERPRETATION
p : [Equation] { s } return x s ::
[Equation]
| s
[Equation]
s | x := e | if ( e> 0) { s } else { s } e ::= x | v | b ( e,e ) | u ( e ) b ::=ADD | MUL u ::=
NEG | SIN | EXP | LOG1P x ::= set of variable names v ::= set of ﬂoating point values
Figure 3 presents the syntax of T URACO , a loop-free IMP-like language (Winskel, 1993).
A T URACO program p takes a list of inputs x , executes a top-level statement s , and returns a single variable x .
Statements s are skips, sequences, assignments, or if statements.
Expressions e are variables x , values v , binary operations b , or unary operations u .
T URACO supports analytic operations (e.g., NEG , SIN , EXP ), including those which can be represented by a power series within a given domain: LOG1P computes log(1+ x ) for x ∈ ( − 1 , 1] .
We restrict the supported operations to those required to implement the case study in Section 5.
Appendix D.1 presents the full set of semantics for T URACO
. 4. C OMPLEXITY A NALYSIS
We now present a program analysis that gives an upper bound on the complexity of traces of T URACO programs.
This section presents a core set of rules; Appendix D.2 presents the full analysis.
The analysis uses three core concepts: a complexity interpretation of expressions to calculate an upper bound on the tilde of expressions based on the calculus in Equation (4), a dual-number execution (Wengert, 1964; Griewank & Walther, 2008) to calculate the derivative of the upper bound on the tilde (to compute the ﬁrst term in Equation (5)), and a path analysis which splits the program by paths to compute the complexity of each trace.
Figure 4 presents the big-step relations used to calculate an upper bound on the complexity for a subset of expressions.
The relation (cid:104) ˜ σ, [ e ] (cid:105) ˜ ⇓ (˜ v, ˜ v (cid:48) ) says that under the variable complexity mapping ˜ σ (mapping variables to tuples with their respective tilde and tilde derivative), the expression e has ˜ e ≤ ˜ v and ˜ e (cid:48) ≤ ˜ v (cid:48) .
Note that the rule for MUL uses the upper bound for multiplication in Equation (4).
Figure 5 presents the complexity relation for if statements.
For statements, (cid:104) ˜ σ
,s (cid:105) ˜ ⇓ ˜Σ says that under the variable complexity mapping ˜ σ , the statement s results in a set of paths with complexity mappings ˜Σ( p ) for path p .
We use a period to denote string concatenation (e.g., l .p to prepend p with l ).
[Figure] [Figure]
To calculate the complexity ζ as deﬁned in Section 3.3 (for a given input β parameter
[Figure]
which represents the scale of the input data) of each path of a program
[Figure]
we use the statement relation to calculate ˜ f (cid:48) ( β ) and ˜ f (0) .
Appendix D.3 presents this rule, along with the theorem that the complexity calculated by this analysis is an upper bound on the complexity as deﬁned in Section 3.
[Figure]
R ENDERER D EMONSTRATION
In this section we present a case study of our optimal sampling results and complexity analysis.
The program under study is a demonstration 3D renderer (Lettier, 2019), such as forms the core of a graphics rendering pipeline for a movie or 3D game engine (Christensen et al., 2018; Tatarchuk, 2006).
Figures 6a and 6b show scenes that the renderer generates.
We demonstrate that the sampling and analysis techniques in Sections 3 and 4 consistently result in more accurate surrogates than those trained using baseline distributions (the frequency distribution of paths and the uniform distribution).
Compared to training surrogates on the frequency distribution of paths, optimal path sampling decreases error by 15% .
Compared to training on the uniform distribution of paths, optimal path sampling decreases error by 47% .
These improvements in error correspond to perceptual improvements in the generated images, as shown in Figures 6c to 6e.
5. P ROGRAM U NDER S TUDY
The full renderer program is a 2750 lines-of-code C++ program, which invokes 38 different GLSL shader programs totaling 2446 lines of code.
We learn a surrogate of a section of one core shader, totaling 60 lines of code.
5 Figure 18 in Appendix E presents the code under study.
Input-output speciﬁcation.
This program is a shader which assigns colors to pixels in the image based on the scene geometry, materials, lights, and other properties.
The program is called for each pixel that is rendered in the image.
Each invocation the of program takes as input a set of 11 ﬁxed-size vectors, totaling 35 inputs.
The program returns as output a set of 2 ﬁxed-size vectors, totalingoutputs.
These outputs are two RGBA colors, the ﬁrst representing the base color of the pixel, and the second representing the color and intensity of a specular map at that pixel.
Scenes and datasets.
We evaluate the renderer on four different scenes, which we combine into nine different datasets.
Figures 6a and 6b present two of the four different scenes under consideration; the four scenes are all combinations of views from the front and top, during the day and night.
We combine these scenes into nine datasets: a dataset with each scene, a dataset combining each scene from each angle (front day and front night, top day and top night), a dataset combining each scene from each time of day, and a dataset combining all scenes.
Figure 14 in Appendix E presents the full set of scenes under study.
[Figure]
[Table] [Table]
Paths.
The program is a conjunction of 48 different paths, 9 of which are exercised by the renderer.
The top part of Table 1 presents statistics about the paths under study, showing the identiﬁer (a trace of l and r characters denoting which branch of each if statement the path takes), the lines of code in the corresponding trace, and the complexity of the corresponding trace according to the analysis in
Section 4.2.
The paths are broken up into a path for rendering smoke particles from the chimney, water particles in the river, and the solids of the ground and house.
Each set of paths is duplicated for twilight, nighttime, and daytime.
Within each time of day, the smoke paths are the least complex, followed by water then solids.
Across time, twilight paths the least complex, followed by daytime then nighttime.
Table 1 also presents the observed and optimal distributions of paths for each dataset.
In general, the twilight paths are rarer than the nighttime paths, which are rarer than the daytime paths: this is because data collection for the nighttime scenes extends through twilight and into the morning.
For all datasets, the smoke paths are rarer than the water paths, which are in turn rarer than the solids paths; this is purely due to the number of points observed for each scene.
Appendix E.1 presents code, statistics, and visualizations of all paths in the scene.
5. S URROGATE T RAINING AND D EPLOYMENT M ETHODOLOGY
To create and deploy a surrogate of the renderer, we train a surrogate of each path, then create a stratiﬁed surrogate which branches on the set of path conditions and applies the corresponding surrogate.
Our goal is to compare the theoretical and empirical errors achieved by training on the optimal sampling distribution against those of baseline distributions.
We compare the approaches across different training datasets, different total numbers of training data points, and evaluating across different evaluation sets, all with multiple trials.
Full methodological details are presented in Appendix E.2.
5. S URROGATE E RRORS
Table 2 presents the geometric mean change in error of using optimal sampling compared to each baseline, on each dataset.
Across most datasets, optimal path sampling results in lower error than both frequency-based path sampling and uniform path sampling.
On datasets with few paths (frontday) and in which all paths are well represented (minimum 5% frequency), the gap is minimal, and frequency-based path sampling matches or outperforms optimal path sampling.
On datasets with more and rarer paths (top-night), the gap widens and optimal path sampling outperforms frequency-based path sampling.
On all datasets, optimal path sampling outperforms uniform path sampling.
5. V ISUALIZATION
Figure 6 presents the renderings generated by the surrogates for the front-day and the top-night scene.
These budgets correspond to the smallest budget that lead to a validation error less than 2% , which was qualitatively chosen as a threshold around which surrogate renders converge on the ground truth.
The top row shows the front day scene using surrogates trained on the front day dataset.
In this scene, the optimal and frequency-based surrogates result in similarly accurate renders, with the primary difference being that the frequency-based surrogate is overall more red, which is most notable in the windows and the front face of the house.
This similarity is expected given the similar errors observed in
Table 2.
Uniform sampling results in an inaccurate render, as expected given its high error.
The bottom row show the top night scene using surrogates trained on the dataset combining all scenes.
In this scene the optimal surrogate has the most accurate render, as expected given the errors observed in Table 2.
The frequency-trained surrogate colors everything (especially the roof and water) slightly more pink.
The uniform-trained surrogate colors everything signiﬁcantly more green.
In sum, the error improvements in Table 2 result in visual improvements in the generated renders.
C ONCLUSION
We present an optimal approach to allocating samples among strata to train stratiﬁed neural network surrogates of stratiﬁed functions.
We also present a programming language, T URACO , in which all programs are learnable stratiﬁed functions and a program analysis to determine the complexity of learning surrogates of those programs.
Our results take a step towards a cohesive, end-to-end methodology for programming using surrogates of programs.
A R ELATED W ORK
In this appendix we survey related work for each contribution.
A. O PTIMAL S TRATIFIED S AMPLING
Optimal stratiﬁed sampling is a classic area in statistics (Thompson, 2012).
Most work in this domain focuses on optimal parameter estimation, and uses stratiﬁed sampling to reduce the variance of estimates by ensuring sufﬁcient independent samples are taken from each subpopulation.
Our approach is novel in the assumptions we make for training stratiﬁed surrogates of programs, and in the speciﬁc sample complexity bounds we base our results on.
Santner et al.
(2018) survey sampling techniques for computer experiments.
Chapter 5.2.3 discusses stratiﬁed random sampling in particular, showing optimality criteria for sampling for unbiased estimators.
These approaches are generic for minimizing the variance of estimators, and do not consider speciﬁcally training a neural network.
These approaches also do not consider the different complexity of different strata.
Cortes et al.
(2019) present an active learning approach for learning in the regime where the input space is partitioned into separate regions (strata, using our terminology) and a separate hypothesis (surrogate) is trained of each, and derive a similar optimal of allocation of data points.
This approach has two primary differences from our approach.
First, it is an active learning approach that determines whether or not to query a label of a given data point for an input stream of data points, whereas our approach operates ofﬂine.
Second, it assumes a different form for sample complexity and derives correspondingly different optimal sampling bounds than ours.
The deﬁnition of complexity ( ζ in our formalism) that
Cortes et al.
use is a function of the number of hypotheses in the hypothesis class, the total number of data points used, and the number of data points for a given stratum that have been queried thus far.
It is not a function of any complexity metric of the function being learned.
Cortes et al. ’s approach is thus a better ﬁt when learning stratiﬁed functions of unknown complexity (i.e., non-analytic functions) using a ﬁnitely sized hypothesis class (not a neural network), and is targeted at the online setting when given a sampler of the overall data distribution but not one for each stratum. Our approach is a better ﬁt when learning stratiﬁed analytic functions with neural networks, and is targeted at the ofﬂine setting when given a sampler for each stratum. A. S AMPLE C OMPLEXITY P ROGRAM A NALYSIS
Program analysis is a broad set of techniques to determine properties of programs (Nielson et al., 1999;
Cousot&Cousot,1977).
OuranalysisinSection4.2isanovelnonstandardinterpretationcalculatingthe tilde, combined with a standard implementation of forward-mode automatic differentiation (Wengert, 1964; Griewank & Walther, 2008) and a standard symbolic execution which executes all paths in the program (King, 1976; Cadar et al., 2008).
Bao et al.
(2012) present a program analysis that decomposes programs into continuous regions, with the goal of characterizing the sensitivity of each continuous region to input noise.
This analysis computes a different notion of complexity than ours, and does not represent the sample complexity of learning a surrogate of each region.
Hoffmann & Hofmann (2010) present a program analysis that calculates the algorithmic complexity of a program.
This complexity again does not lead to bounds on the sample complexity of learning a surrogate of the program.
B A SSUMPTIONS
Our contributions in Sections 3 and 4 make assumptions about the programs under study, the functions that those programs denote, and the surrogate training algorithms.
Here we document these assumptions and note possible failure modes for our techniques.
Assumptions imported from prior work.
Our sample complexity results are subject to all assumptions from the prior work that gives the sample complexity bounds for neural networks that we use (Agarwala et al., 2021).
These sample complexity bounds only apply to analytic functions.
They further assume that inputs come from the unit sphere; this does not match many practical applications, including those in Section 5.
Finally, these sample complexity bounds assume that the neural network under study is a 2-layer, inﬁnitely wide neural network trained with SGD with an inﬁnitely small step size, using a 1-Lipschitz loss function.
Despite these assumptions, Agarwala et al.
(2021, Appendix B.2) empirically verify that the sample complexity bounds hold.
We also show in Section 2 that the theoretical sample complexity bounds correlate with empirical sample complexity results on three functions.
Optimal sampling.
The ﬁrst assumption is that we know the distribution of inputs ahead of time
D ( x ) , both in terms of the distribution of strata D ( s i ) and the distribution of inputs within a given stratum D ( x | s i ) .
The second assumption is that optimizing the upper bound of the per-stratum loss results in a reasonable optimum for the combined surrogate.
The third is the technical assumption we make that δ i = δc ; this is not guaranteed to be optimal.
Program analysis.
The main assumption here is that Agarwala et al.
’s algebra on tilde functions results in a sufﬁciently precise upper bound on the tilde.
This is not always the case – for instance, this assumption says that ADD(x, NEG(x)) has a positive tilde (since the tilde of NEG(x) is equal to the tilde of x ), whereas it should have a tilde of 0.
Stratiﬁed surrogates.
We provide sample complexity bounds for constructing stratiﬁed surrogates, assuming that for a given program every path is a different function.
This assumes both that it is tractable to compute which stratum a given input resides in before applying the surrogate.
This also assumes that there are a tractable number of paths, which excludes programs with a large number of if statements or loops.
C P ROOF OF O PTIMAL S AMPLING
This appendix provides the proof for the optimal sampling result in Section 3.
For convenience we duplicate the deﬁnitions of learnability.
A learnable function is one such that:
[Equation]
Following Arora et al.
(2019) and Agarwala et al.
(2021) we study functions and learning algorithms for which the relationship between n , (cid:15) , and δ in Equation (1) is:
[Equation]
The optimization problem of our optimal sampling approach is:
[Equation]
Theorem 3.1.
With δ i = δc , the upper bound of Equation (3) is minimized at:
[Equation]
Proof.
The task is to ﬁnd n i for each surrogate
ˆ f i that minimize the expectation of error and such that (cid:80) i n i ≤ n .
We ﬁrst separate out the expectation by stratum, and upper bound the error using the sample complexity assuming for the surrogates:
[Equation]
Plugging in δ i and writing out the Lagrangian [Equation] we have:
[Equation]
Because the objective is convex in n i , this is the optimal solution.
D F ULL T URACO S EMANTICS
This appendix presents the full semantics and analysis for T URACO .
Figure 7 re-presents T URACO .
D. N ORMAL E VALUATION S EMANTICS
Figure 8 presents the big-step evaluation relation for expressions in T URACO .
The expression relation (cid:104) σ,e (cid:105) ⇓ v says that under variable store σ (assigning values to all variables in e ), the expression e evaluates to value v .
[Equation]
s }
[Equation]
skip
[Equation]
x :
[Equation]
set of variable names v ::= set of ﬂoating point values
[Figure]
[Figure] [Figure]
Figure 9 presents the big-step evaluation relation for statements in T URACO .
The statement relation (cid:104) σ,s (cid:105)⇓ σ (cid:48) says that under variable store σ , the statement s evaluates to a new variable store σ (cid:48) .
Figure 10 presents the big-step evaluation relation for T URACO programs.
The program relation (cid:104) σ, fun( x 0 ,x 1 ...,x n ) { s } return x (cid:105)⇓ v says that under variable store σ (assigning values to all x i ) the program evaluates to value v .
D. C OMPLEXITY A NALYSIS
Figure 11 presents the full complexity relation used to calculate the complexity for all expressions.
The relation (cid:104) ˜ σ, [ e ] (cid:105) ˜ ⇓ (˜ v, ˜ v (cid:48) ) says that under the variable complexity mapping ˜ σ (mapping variables to tuples with their respective tilde and tilde derivative), the expression e has ˜ e ≤ ˜ v and ˜ e (cid:48) ≤ ˜ v (cid:48) .
Figure 12 presents the complexity relation for statements.
For statements, the relation (cid:104) ˜ σ,s (cid:105) ˜
⇓ ˜Σ says that under the variable complexity mapping ˜ σ , the statement s results in a set of paths with variable complexity mappings ˜Σ( p ) for path p .
Figure 13 presents the complexity ζ (for a given input β parameter) of each path of a program.
D. A NALYSIS C ORRECTNESS
This appendix presents the statement of correctness for the T URACO complexity analysis.
Theorem 1.
Given:
[Equation]
then for all paths p , ζ β ( f ) ≤ ˜ v p , where f is the function denoted by the standard execution of the path p through s .
Proof.
The proof is by induction using Fact 1 (p. 17) and Lemma 7 (p. 22) in Agarwala et al. [Table] [Figure] [Figure]
[Figure]
{+(cid:48) p, 0 ) ∀ p. ˜Σ β ( p )( x )=(˜ v p,β , ˜ v (cid:48) p,β ) ζ β [fun( {{ x i } ) { s } return x ] ≤ (cid:110) p (cid:55)→ (cid:0) ˜ v p, 
+}
+˜ v (cid:48) p,β (cid:1) 2 (cid:12)(cid:12)(cid:12) p ∈ ˜Σ 0 (cid:111)
E E XTENDED R ENDERER D EMONSTRATION
This appendix provides further details on the evaluation on the renderer program in Section 5.
Figure 18 presents the full code for the renderer program under study.
7 Figure 14 presents the full set of scenes used in our evaluation.
This program is a good candidate to train a surrogate of for several reasons.
First, it is an approximable program: as long as the outputs of a surrogate of the program are sufﬁciently close to the ground-truth outputs, the generated image will be perceptually indistinguishable.
Second, its paths are all determined by uniform input variables, variables in GLSL that are constant across each invocation of the shader.
This means that relative to the cost of executing the program, it is cheap to determine which path a given input induces in the program (and thus which surrogate to apply).
Third, its execution environment is well suited to be replaced with a neural network, since the original program itself performs batch processing on a GPU.
E. P ATHS
Table 3 presents statistics about the entire set of paths paths under study.
Figure 16 presents a side-by-side comparison of traces of two paths: lrrlrl (twilight smoke), the least complex path in the dataset, and lrrrrr (nighttime solids), the most complex path in the dataset. [Figure] [Figure]
Two factors lead to the twilight smoke path being less complex than the nighttime solids path.
First, the twilight smoke path computes one lighting factor from a constant (Figure 16a, line 14) rather than from a function of the sun’s position (Figure 16b, lines 13-14).
Second, the twilight smoke path does not have any specular output (Figure 16a, line 18), unlike the nighttime solid output (Figure 16b, line 18).
Figure 15 shows side-by-side comparisons of the three classes of paths: water, smoke, and solids.
In each of these images, one path returns red for all pixels, while the other paths return black for all pixels.
The base scene is the front daytime scene in Figure 19a. [Table] [Figure]
fun ( . . . )
{ 2 sunPosition := sin (mul( sunPosition [0] , (cid:44) → mul( pi , 0.0055556) ) ) ...
15 emission:=mul(emission,0.1) 16 out0rgb := add ( . . .
, emission ) 17 out0a := diffuseColor [3] 18 out1a := diffuseColor [3] 19 out1rgb:=[0.0,0.0,0.0] [Figure] E. A RCHITECTURE AND H YPERPARAMETER D ETAILS
For the design of each surrogate, we use a simple MLP architecture with a single hidden layer ofunits and a ReLU activations.
This architecture matches that of Agarwala et al.
(2021), except using 512 rather than 1000 hidden units (we found that the accuracy of each path surrogate plateaued byhidden units).
We train the surrogate using the Adam optimizer with a learning rate of 0 . 0001 and a batch size of 256 for 50,000 steps.
We run 5 trials of all experiments, and report the mean error.
E. F ULL R ESULTS
Figure 17 presents the error of surrogates on each dataset.
Each plot shows the error for a different dataset.
Each plot has three different lines, respectively showing the error of each surrogate training distribution (optimal, frequency-based, and uniform).
Each x axis is the total training data budget.
Each y axis is the error of the resulting stratiﬁed surrogate.
[Figure] Figures 19 and 20 present larger versions of the renders in Section 5.4 for ease of comparison.
The training budgets for these surrogates are 61 samples for the front-day surrogate and 206 samples for the top-night surrogate.
fun ( rimLight [4] , isCelShadingEnabled [2] , sunPosition [2] , gamma[2] , (cid:44) → worldNormal [3] , ssao [3] , diffuseColor [4] , diffuse [4] , specular [4] , (cid:44) → emission [3] , isWater [2] , i s P a r t
i c l e [2 ]) {sunPosition = sin ( mul ( sunPosition [0] , mul ( pi , 0.0055556) ) )sunMixFactor
= add (0.5 , neg ( mul ( sunPosition , 0.5) ) )
ambientCoolBase = exp ( mul ( log ( add ((0.302 , 0.451 , 0.471) , −1) ) , (cid:44) → gamma [0] ) ) ambientWarmBase = exp ( mul ( log ( add ((0.765 , 0.573 , 0.400) , −1) ) , (cid:44) → gamma [0] ) )
i f (0.5 > sunMixFactor )
{ambientCool = mul ( ambientCoolBase , 0.5)ambientWarm =
mul ( ambientWarmBase , 0.5)} else {ambientCool = mul ( ambientCoolBase , sunMixFactor )
ambientWarm =
mul ( ambientWarmBase , sunMixFactor )} i f (0 > sunMixFactor )
{skyLight = ambientCoolgroundLight =
ambientWarm} else {i f ( sunMixFactor > 1) {skyLight = ambientWarmgroundLight
= ambientCool} else {skyLight = add ( mul ( ambientCool , add (1 , neg ( sunMixFactor ) ) ) , (cid:44) →
mul ( ambientWarm , sunMixFactor ) )
groundLight = add ( mul ( ambientWarm , add (1 , neg ( sunMixFactor ) ) ) , (cid:44) →
mul ( ambientCool , sunMixFactor ) ) } 26 } worldNormalMixFactor = mul (0.5 , add (1.0 , worldNormal [2]) )
ambientLight = add ( mul ( groundLight , add (1 , neg ( worldNormalMixFactor ) ) ) , (cid:44) →
mul ( skyLight , worldNormalMixFactor ) ) ambient =
mul ( ambientLight , mul (( diffuseColor [0] , diffuseColor [1] , (cid:44) → diffuseColor [2 ]) , ssao ) )
i f (0.00316228 > sunPosition ) {emission = mul ( emission , 0.1)} else {sunPositionPow = exp ( mul ( log ( add ( sunPosition , −1) ) , 0.4) )emission =
mul ( emission , sunPositionPow )}
out0rgb = add ( add (( ambient [0] , ambient [1] , ambient [2]) , ( diffuse [0] , (cid:44) → diffuse [1] , diffuse [ 2]) ) , add (( rimLight [0] ,
rimLight [1] , (cid:44) →
rimLight [2] ) , emission ) )
i f
( isWater [0] > 0) {out0a = 0} else {out0a = diffuseColor [3]}out1a = diffuseColor [3] i f
( i s P a r t i c
l e [0] > 0) {out1rgb = (0 ,0 ,0)} else {out1rgb = ( specular [0] , specular [1] , specular [2])} 52 } return ( out0rgb , out0a , out1rgb , out1a )
[Figure] [Figure]