Question: Weight-sharing—the simultaneous optimization of multiple neural networks using the same parameters—has emerged as a key component of state-of-the-art neural architecture search. However, its success is poorly understood and often found to be surprising. We argue that, rather than just being an optimization trick, the weight-sharing approach is induced by the relaxation of a structured hypothesis space, and introduces new algorithmic and theoretical challenges as well as applications beyond neural architecture search. Algorithmically, we show how the geometry of ERM for weight-sharing requires greater care when designing gradient- based minimization methods and apply tools from non-convex non-Euclidean optimization to give general-purpose algorithms that adapt to the underlying structure. We further analyze the learning-theoretic behavior of the bilevel optimization solved by practical weight-sharing methods. Next, using kernel configuration and NLP feature selection as case studies, we demonstrate how weight-sharing applies to the architecture search generalization of NAS and effectively optimizes the ing bilevel objective. Finally, we use our optimization analysis to develop a simple exponentiated gradient method for NAS that aligns with the underlying optimization geometry and matches state-of-the-art approaches on CIFAR-10. Weight-sharing neural architecture search (NAS) methods have achieved state-of-the-art performance while requiring computation training of just a single shared-weights network (; ; . However, weight-sharing remains poorly understood. In this work, we present a novel perspective on weight-sharing NAS motivated by the key observation that these methods subsume the architecture hyperparameters as another set of learned parameters of the shared-weights network, in effect extending the hypothesis class. An important ramification of this insight is that weight-sharing is not NAS-specific and can be used to tune hyperparameters corresponding to parameterized feature maps of the input data. We refer this larger subset of hyperparameter optimization problems as architecture search, and we study the following two questions associated with weight-sharing applied to the architecture search problem: 1. How can we efficiently optimize the objective induced by applying weight sharing to architecture search, namely minimizing empirical risk in the joint space of model and architecture parameters? For large structured search spaces that preclude brute force search, a natural approach to architecture search with weight-sharing is to use gradient-based methods to minimize the empirical risk over a continuous relaxation of the discrete space . Although this has allowed NAS researchers to apply their preferred optimizers to determine architecture weights, it is far from clear that the success of established methods for unconstrained optimization in training neural networks will naturally extend to these constrained and often non-Euclidean environments. As we foresee that architecture search spaces will continue to become more complex and multi-faceted, we argue for and develop a more principled, geometry-aware formulation of the optimization problem. Drawing upon the mirror descent meta-algorithm and successive convex approximation, we give non-asymptotic stationary-point convergence guarantees for the empirical risk minimization (ERM) objective associated with weight-sharing via algorithms that simultaneously connect to the underlying problem structure and handle the alternating-block nature of the architecture search. Our guarantees inform the design of gradient-based weight-sharing methods by explicitly quantifying the impact of optimizing in the right geometry on convergence rates. 2. What are the generalization benefits of solving a bilevel optimization for the architecture search problem commonly considered in practice? At its core, the goal of architecture search is to find a configuration that achieves good generalization performance. Consequently, a bilevel objective that optimizes the architecture weights using a separate validation loss is commonly used in practice in lieu of the ERM objective naturally induced by weight sharing . The learning aspects of this approach have generally been studied in settings with much stronger control over the model complexity . We provide generalization guarantees for this objective over structured hypothesis spaces associated with a finite set of architectures; this leads to meaningful bounds for simple feature map selection problems as well as insightful for the NAS problem that depend on the size of the space of global optima. To validate our theoretical , we conduct empirical studies of weight-sharing in two settings: shallow feature map selection, i.e., tuning the hyperparameters of kernel classification and NLP featurization pipelines, and CNN neural architecture search. In we demonstrate that weightsharing efficiently optimizes the bilevel objective and achieves low generalization error with respect to the best architecture setting. For, motivated by insights from our convergence analysis, we develop a simple exponentiated gradient version of DARTS called EDARTS that better exploits the geometry of the optimization problem. We evaluate EDARTS on the design of CNN architectures for CIFAR-10 and demonstrate that EDARTS finds better architectures than DARTS in less than half the time. We also achieve very competitive relative to state-of-the-art architectures when using an extended evaluation routine. Related Work: Our work on optimization for weight-sharing benefits from the literature on firstorder stochastic optimization and in particular the mirror descent framework . Specifically, we use successive convex approximation to show convergence of alternating minimization and derive geometry-dependent rates comparable to existing work on non-convex stochastic mirror descent . Our generalizes to the constrained, nonEuclidean, and multi-block setting an approach of for obtaining non-convex convergence from strongly convex minimization, which may be of independent interest. Previous optimization for NAS have generally only shown bounds on auxiliary quantities such as regret that are not well-connected to the learning objective ) or have only given monotonic improvement or asymptotic guarantees . However, due to the generality of mirror descent, the approaches in the middle three papers can be seen as special cases of our analysis. Finally, our analysis of the properties of the bilevel optimization is related to work on model selection , but does not consider the configuration parameters as explicit controls on the model complexity. Our learning are broadly related to hyperparameter optimization, although most work focuses on algorithmic and not statistical questions . In this section, we formalize the weight-sharing learning problem, relate it to traditional ERM, and provide examples for the case of NAS and feature map selection that we use for the rest of the paper. Our main observation is that weight-sharing for architecture search extends the hypothesis space to be further parameterized by a finite set of configurations C. Formally, we have a structured hypothesis space H(C, W) = {h w ∈ H(C, W) with low population error w (x), y) for loss: Y × Y → R. Hence, we can apply ERM as usual, with optional regularization, to select a hypothesis from the extended hypothesis space; in fact this is done by some NAS methods (e.g.,). The learning algorithm is then min w∈W,c∈C for block specific regularizers R W and R C. Note that in the absence of weight-sharing, we would need to learn a separate hypothesis h (c) wc for each hypothesis subclass H c. Although a brute force approach to selecting a hypothesis from H(C, W) via ERM would in effect require this as well, our subsequent examples demonstrate how the weight-sharing construct allows us to apply more efficient gradient-based optimization approaches, which we study in Section 3. Feature Map Selection: In this setting, the structure is induced by a set of feature transformations C = {φ i : X → R n for i = 1, . . ., k}, so the hypothesis space is {f w (φ i (·)): w ∈ W, φ i ∈ C} for some W ⊂ R d. Examples of feature map selection problems include tuning kernel hyperparameters for kernel ridge classification and tuning NLP featurization pipelines for text classification. In these cases f w is a linear mapping f w (·) = w, · and W ⊂ R n. Neural Architecture Search: Weight-sharing methods almost exclusively use micro cell-based search spaces for their tractability and additional structure (; . These search spaces can be represented as directed acyclic graphs (DAGs) with a set of ordered nodes N and edges E. Each node x (i) in the DAG is a feature representation and each edge o (i,j) is an operation on the feature of node j passed to node i and aggregated with other inputs to form x (j), with the restriction that a given node j can only receive edges from prior nodes as input. Hence, the feature at a given node i is. Search spaces are then specified by the number of nodes, the number of edges per node, and the set of operations O that can be applied at each edge. In this case the structure C ⊂ {0, 1} |E||O| of the hypothesis space is the set of all valid architectures for this DAG encoded by edge and operation decisions. Treating both weights w ∈ W and architecture decision c ∈ C as parameters, weight-sharing methods train a single shared-weights network h (c) w: X → Y encompassing all possible functions within the search space. Therefore, the sharedweights network includes all possible edges between nodes and all possible operations per edges. In addition to the weights w ∈ W corresponding to all the operations, the shared-weights network also takes architecture weights c ∈ C as input, where c indicates the weight given to operation o on edge (i, j) so that the feature of a given node i is Gradient-based weight-sharing methods apply continuous relaxations to the architecture search space in order to compute gradients. Some methods like DARTS and its variants (; ; ; ; relax the search space by considering a mixture of operations per edge and then discretize to a valid architecture in the search space. With the mixture relaxation, we replace all c ∈ {0, 1} |E||O| in the above expressions by continuous counterparts θ ∈ |E||O|, with the constraint that o∈O θ (i,j) o = 1, i.e., the architecture weights for operations on each edge sum to 1. Other methods like SNAS, ASNG-NAS , and ProxylessNAS assume a parameterized distribution p θ from which architectures are sampled. By substituting continuous parameters θ ∈ Θ in for discrete parameters c ∈ C, we are able to use gradient-based methods to optimize. We address the question of how to effectively use gradient optimization for weight-sharing in the next section. While continuous relaxation enables state-of-the-art , architecture search remains expensive and noisy, with state-of-the-art mixture methods requiring second-order computations and probabilistic methods suffering from high variance policy gradients. Moreover, while the use of SGD to optimize network weights is a well-tested approach, architecture weights typically lie in constrained, non-Euclidean geometries in which other algorithms may be more appropriate. Recognizing this, several efforts have attempted to derive a better optimization schemes; however, the associated guarantees for most of them hold for auxiliary objectives, such as regret of local configuration decisions, that are not connected to the optimization objective and ignore the two-block nature of the problem ). do consider an alternating descent method for the training objective, their are asymptotic and certainly do not indicate any finite-time convergence. In this section we address these difficulties by showing that the mirror descent (MD) framework is the right tool for designing algorithms in the block optimization problems that occur in architecture search. We describe how such geometry-aware gradient algorithms lead to faster stationary-point convergence; as we will show in Section 5, this yields simple, principled, and effective algorithms for large-scale multi-geometry problems such as NAS. Algorithm 1: Two geometry-aware optimization algorithms for multi-block optimization for a β-strongly-smooth function over if SBMD Problem Geometry: We relax optimization problems of the form to the problem of minimizing a function f: X → R over a convex product space X i consisting of blocks i, each with an associated norm · (i). For example, in typical NAS we can set X 1 to be the set-product of |E| simplices over the operations O and the associated norm to be · 1 while X 2 = W ⊂ R d is the space of network weights and the associated norm is · 2. To each block i we further associate a distance-generating function (DGF) ω i: X → R that is 1-strongly-convex w.r.t. . For example, in the Euclidean case using ω 2 (·) = yields the usual squared Euclidean distance; over the probability simplex we often use the entropy ω 1 (·) = ·, log(·), which is 1-strongly-convex w.r.t. · 1 and for which D ω1 is the KL-divergence. Given an unbiased gradient estimate g(x t, ζ) = E ζ ∇f (x t), the (single-block) stochastic MD step is for some learning rate η > 0. In the Euclidean setting this reduces to SGD, while with the entropic regularizer the iteration becomes equivalent to exponentiated gradient. Key to the guarantees for MD is the fact that the dual norm of the problem geometry is used to measure the second moment of the gradient; if every coordinate of g is bounded a.s. by σ, then under Euclidean geometry the dependence is on E g(x) 2 2, * = E g(x) 2 2 ≤ σ 2 d while using entropic regularization it is on. Thus in such constrained 1 geometries mirror descent can yield dimension-free convergence guarantees. While in this paper we focus on the benefit in this simple case, the MD meta-algorithm can be used for many other geometries of interest in architecture search, such as for optimization over positive-definite matrices . Algorithms and Guarantees: We propose two methods for the above multi-block optimization problems: stochastic block mirror descent (SBMD) and alternating successive convex approximation (ASCA). At each step, both schemes pick a random coordinate i to update; SBMD then performs a mirror descent update similar to but with a batched gradient, while ASCA optimizes a strongly-convex surrogate function using a user-specified solver. Note that both methods require that f is β-strongly-smooth w.r.t. each block's norm · (i) to achieve convergence guarantees, a standard assumption in stochastic non-convex optimization . This condition holds for the architecture search under certain limitations, such as a restriction to smooth activations. In the supplement we also show that in the single-block case ASCA converges under the more general relative-weak-convexity criterion . We first discuss SBMD, for which non-convex convergence guarantees were shown by; this algorithm is the one we implement in Section 5 for NAS. A first issue is how to measure stationarity in constrained, non-Euclidean geometries. In the single-block setting we can set a smoothness-dependent constant λ > 0 and measure how far the proximal gradient operator prox∇ λ (x) = arg min u∈X λ ∇f (x), u + D ω (u||x) is from a fixed point, which yields the projected gradient measure. Notably, in the unconstrained Euclidean case this yields the standard stationarity measure ∇f (x) 2 2. For the multi-block case we replace an ε-stationary point of f w.r.t. the projected gradient if G λ (x) ≤ ε. In the non-Euclidean case the norm of the projected gradient measures how far we are from satisfying a first-order optimality condition, namely how far the negative gradient is from being in the normal cone of f to the set X (, Proposition 4.1). For this condition, Algorithm 1 has the following guarantee: Theorem 3.1 . If f is β-strongly-smooth, F = f (x ) − min u∈X f (u), and oracle calls to reach an ε-stationary-point x ∈ X as measured by G 1 We next provide a guarantee for ASCA in the form of a reduction to strongly-convex optimization algorithms. This is accomplished by the construction of the solving of a surrogate function at each iteration, which in the Euclidean case is effectively adding weight-decay. ASCA is useful in the case when efficient strongly-convex solvers are available for some or all of the blocks; for example, this is frequently the case for feature map selection problems such as our kernel approximation examples, which employ 2 -regularized ERM in the inner optimization. Taking inspiration from , for ASCA we analyze a stronger notion of stationarity that upper-bounds G λ (x) 2, which we term the projected stationarity measure:. Note that this is not the same measure used by , although in the appendix we show that in the single-block case our holds for their notion also. Theorem 3.2. If f is β-strongly-smooth and where ε > 0 is the solver tolerance and the expectation is over the randomness of the algorithm and associated oracles. Proof Summary. The proof generalizes a of and is in Appendix A. Thus if we have solvers that return approximate optima of strongly-convex functions on each block then we can converge to a stationary point of the original function. The convergence rate will depend on the solver used; for concreteness we give a specification for stochastic mirror descent. Corollary 3.1. Under the conditions of Theorem 3.1, if on i ASCA uses the Epoch-GD method of ωi oracle calls suffice to reach an ε-stationary-point x ∈ X as measured by E ∆ 1 This oracle complexity matches that of Theorem 3.1 apart from the extra β 2 L 2 ωi term due to the surrogate function, which we show below is in practice a minor term that does not obviate the benefit of geometry-awareness. On the other hand, ASCA is much more general, allowing for many different algorithms on individual blocks; for example, many popular neural optimizers such as Adam have variants with strongly-convex guarantees . The Benefit of Geometry-Aware Optimization: We conclude this section with a formalization of how convergence of gradient-based architecture-search algorithms can benefit from this optimization strategy. Recalling from our example, we have the architecture parameter space X 1 consisting of |E| simplices over |O| variables equipped with the 1-norm and the shared weight space X 2 = W equipped with the Euclidean norm. We suppose that the stochastic gradients along each block i have coordinates bounded a.s. by σ i > 0. Then if we run SBMD using SGD to optimize the shared weights and exponentiated gradient to update the architecture parameters, Theorem 3.1 implies that we reach an ε-stationary point in O stochastic gradient computations. The main benefit here is that the first term in the numerator is not σ 2 1 |E||O|, which would be the case if we used SGD; this improvement is critical as the noise σ 2 1 of the architecture gradient can be very high, especially if a policy gradient is used to estimate probabilistic derivatives. In the case of ASCA, we can get similar guarantees assuming the probabilities for each operation are lower-bounded by some small δ > 0 and that the space of shared weights is bounded by B; then the guarantee will be as above except with an additional O(log 2) term (independent of σ 1, σ 2). While for both SBMD and ASCA the σ 2 2 d term from training the architecture remains, this will be incurred even in single-architecture training using SGD. Furthermore, in the case of ASCA it may be improved using adaptive algorithms . In Section 2, we described the weight-sharing hypothesis class H(C, W) as a set of functions nondisjointly partitioned by a set of configurations C sharing weights in W and posed the ERM problem associated with selecting a hypothesis from H(C, W). However, as mentioned in Section 1, the objective solved in practice is a bilevel problem where a separate validation set is used for architecture parameter updates. Formally, the bilevel optimization problem considered is min w∈W,c∈C where T, V ⊂ Z is a pair of training/validation sets sampled i.i.d. from D, the upper objective w (x), y) is the empirical risk over V, and L T (w, c) is some objective induced by T. We intentionally differentiate the two losses since training is often regularized. This setup is closely related to the well-studied problems of model selection and cross-validation. However, a key difference is that the choice of configuration c ∈ C does not necessarily provide any control over the complexity of the hypothesis space; for example, in NAS as it is often unclear how the hypothesis space changes due to the change in one decision. By contrast, the theory of model selection is often directly concerned with control of model complexity. Indeed, in possibly the most common setting the hypothesis classes are nested according to some total order of increasing complexity, forming a structure . This is for example the case in most norm-based regularization schemes. Even in the non-nested case, there is often an explicit tradeoff between parsimony and accuracy . With the configuration parameters in architecture search behaving more like regular model parameters rather than as controls on the model complexity, it becomes reasonable to wonder why most NAS practitioners have used the bilevel formulation. Does the training-validation split exploit the partitioning of the hypothesis space H(C, W) induced by the configurations C? To see when this might be true, we first note that a key aspect of the optima of the bilevel weight-sharing problem is the restriction on the model weights -that they must be in the set arg min w∈W L T (h (c) w ) of the inner objective L T. As we will see, under certain assumptions this can reduce the complexity of the hypothesis space without harming performance. First, for any sample w )} be the version space (, Equation 6) induced by some configuration c and the objective function. Second, let N (F, ε) be the L ∞ -covering-number of a set of functions F at scale ε > 0, i.e. the number of L ∞ balls required to construct an ε-cover of F (, Equation 3 .60). These two quantities let us define a complexity measure over the shared weight hypothesis space: The version entropy is a data-dependent quantification of how much the hypothesis class is restricted by the inner optimization. For finite C, a naive bound shows that Λ(H, ε, T) is bounded by log |C| + max c∈C log N (H c (T), ε), so that the second term measures the worst-case complexity of the global minimizers of L T. In the feature selection problem, L T is usually a strongly-convex loss due to regularization and so all version spaces are singleton sets, making the version entropy log |C|. In the other extreme case of nested model selection the version entropy reduces to the complexity of the version space of the largest model and so may not be informative. However, in practical problems such as NAS an inductive bias is often imposed via constraints on the number of input edges. To bound the excess risk in terms of the version entropy, we first discuss an important assumption that describes cases when we expect the shared weights approach to perform well: Assumption 4.1. There exists a good c * ∈ C, i.e. one satisfying (w *, c w) for some w * ∈ W, such that w.h.p. over the drawing of training set T ∼ D m T at least one of the minima of the optimization induced by c * and T has low excess risk, i.e. w.p. 1 − δ there exists This assumption requires that w.h.p. the inner optimization objective does not exclude all low-risk classifiers for the optimal configuration. Note that it asks nothing of either the other configurations in C, which may be arbitrarily bad, nor of the hypotheses found by the procedure. It does however prevent the case where one knows the optimal configuration but minimizing the provided objective L T does not provide a set of good weights. Note that if the inner optimization is simply ERM over the training set T, i.e. L T = T, then standard learning-theoretic guarantees will give ε * exc (m T, δ) decreasing in m T and increasing at most poly-logarithmically in 1 δ. With this assumption, we can show the following guarantee on solutions to the bilevel optimization. Theorem 4.1. Letĥ be a hypothesis corresponding to the solution of the bilevel optimization. Then under Assumption 4.1 if is B-bounded we have w.p. 1 − 3δ that The first difference is bounded by the version entropy usng the constraint onĥ ∈ H c, the second by optimality ofĥ on V, the third by Hoeffding's inequality, and the last by Assumption 4.1. As shown in the applications below, the significance of this theorem is that a bound on the version entropy guarantees excess risk almost as good as that of the (unknown) optimal configuration without assuming anything about the complexity or behavior of sub-optimal configurations. Feature Map and Kernel Selection: In the feature map selection problem introduced in Section 2, } is a set of feature maps and the inner problem L T is 2 -regularized ERM for linear classification over the ing feature vectors. The bilevel problem is then Due to strong-convexity of L T, each map φ i induces a unique minimizing weight w ∈ W and thus a singleton version space, therefore upper bounding the version entropy by log |C| = log N. Furthermore, for Lipschitz losses and appropriate choice of regularization coefficient, standard for 2 -regularized ERM for linear classification (e.g. In the special case of kernel selection using random Fourier approximation, we can apply associated generalization guarantees (, Theorem 1) to show that we can compete with the optimal RKHS from among those associated with one of the configurations: Corollary 4.2. In feature map selection suppose each map φ ∈ C is associated with a random Fourier feature approximation of a continuous shift-invariant kernel that approximates an RKHS H φ and is the square loss. If the number of features In both cases we are able to get risk bounds almost identical to the excess risk achievable if we knew the optimal configuration beforehand, up to an additional capacity term depending weakly on the number of configurations. This would not be possible with solving the regular ERM objective instead of the bilevel optimization as we would then have to contend with the possibly high complexity of the hypothesis space induced by the worst configuration. Neural Architecture Search: In the case of NAS we do not have a bound on the version entropy, which now depends on all of C. Whether the version space, and thus the complexity, of deep networks is small compared to the number of samples is unclear, although we gather some evidence. number of critical points is exponential only in the number of layers, which would yield a small version entropy. It is conceivable that the quantity may be further bounded by the complexity of solutions explored by the algorithm when optimizing L T ; indeed, we find that shared-weight optimization leads to models with smaller 2 -norm and distance from initialization than from-scratch SGD on a single network (see Appendix D.4). On the other hand, argue, with evidence in restricted settings, that even the most stringent implicit regularization cannot lead to a non-vacuous uniform convergence bound; if true more generally this would imply that the NAS version entropy is quite large. Here we demonstrate how weight-sharing can be used as a tool to speed up general architecture search problems by applying it to two feature map selection problems. We then validate our optimization analysis with a geometry-aware weight-sharing method to design CNN cells for CIFAR-10. Feature Map Selection: Recall that here our configuration space has k feature maps φ i: X → R n with outputs passed to a linear classifier w ∈ R n, which will be the shared weights. We will approximate the bilevel optimization with the inner minimization over 2 -regularized ERM λ w 2 2 + (x,y)∈T (w, φ i (x), y). Our weight-sharing procedure starts with a vector θ ∈ ∆ N encoding a probability distribution p θ over [N] and proceeds as follows: according to (an estimate of) its validation loss (x,y)∈V (w (t), φ i (x), y). Observe the equivalence to probabilistic NAS: at each step the classifier (shared parameter) is updated using random feature maps (architectures) on the training samples. The distribution over them is then updated using estimated validation performance. We consider two schemes for this update of θ (t): exponentiated gradient using the score-function estimate and successive elimination, where we remove a fraction of the feature maps that perform poorly on validation and reassign their probability among the remainder. may be viewed as a softer version of, with halving also having only one hyperparameter (elimination rate) and not two (learning rate, stopping criterion). The first problem we consider is kernel ridge regression over random Fourier features on CIFAR-10. We consider three configuration decisions: data preprocessing, choice of kernel, and bandwidth parameter. This problem was considered by , except they fixed the Gaussian kernel whereas we also consider Laplacian; however, they also select the regularization parameter λ, which weight-sharing does not handle. We also study logistic regression for IMDB sentiment analysis of Bag-of-n-Gram (BonG) featurizations, a standard NLP baseline . Here there are eight configuration decisions: tokenization method, whether to remove stopwords, whether to lowercase, choice of n, whether to binarize features, type of feature weighting, smoothing parameter, and post-processing. As some choices affect the feature dimension we hash the BonGs into a fixed number of bins . To test the performance of weight-sharing for feature map selection, we randomly sample 64 configurations each for CIFAR-10 and IMDB and examine whether the above schemes converge to the optimal choice. The main comparison method here is thus random search, which runs a full sweep over these samples; by contrast successive halving will need to solve 6 = log 2 64 regression problems, while for exponentiated gradient we perform early stopping after five iterations. Note that weight-sharing can do no better than random search in terms of accuracy because they are picking a configuration from a space that random search sweeps over. The goal is to see if it consistently returns a good configuration much faster. As our in Figures 1 and 2 show, successive halving indeed does almost as well as random search in much less time. While exponentiated gradient usually does not recover a near-optimal solution, it does on average return a configuration in the top 10%. We also note the strong benefit of over-parameterization for IMDB -the n-gram vocabulary has size 4 million so the number of bins on the right is much larger than needed to learn in a singleconfiguration setting. Overall, these experiments show that weight-sharing can also be used as a fast way to obtain signal in regular learning algorithm configuration and not just NAS. NAS on CIFAR-10: Recall from Section 3 that when the architecture space consists of |E| simplices of dimension |O|, the convergence rate of exponentiated gradient descent to a stationary point of the objective function is independent of the dimension of the space, while SGD has linear dependence. This motivates our geometry-aware method called Exponentiated-DARTS (EDARTS). EDARTS modifies first-order DARTS in two ways. First, in lieu of the softmax operation used by DARTS on the architecture weights, we use standard normalization so that the weight of operation o on edge (i, j) is u. Second, in lieu of Adam, we use exponentiated gradient to update the architecture weights: c t = c t−1 exp(−η∇ c V (h wt−1 (c t−1)). While EDARTS resembles XNAS, our justification for using exponentiated gradient comes directly from aligning with the optimization geometry of ERM. Additionally, EDARTS only requires two straightforward modifications of first-order DARTS, while XNAS relies on a wipeout subroutine and granular gradient-clipping for each edge operation on the cell and data instance level. 1 1 Our own XNAS implementation informed by correspondence with the authors did not produce competitive . We still compare to the architecture XNAS reported evaluated by the DARTS training routine in Table 1. We evaluate EDARTS on the task of designing a CNN cell for CIFAR-10. We use the standard search space as introduced in DARTS ) for evaluation we use the same three stage process used by DARTS and random search with weight-sharing , with stage 3 considered the'final' . We provide additional experimental details in Appendix D. Table 1 shows the performance of EDARTS relative to both manually designed and NAS-discovered architectures. EDARTS finds an architecture that achieves competitive performance with manually designed architectures which have nearly an order-of-magnitude more parameters. Additionally, not only does EDARTS achieve significantly lower test error than first-order DARTS, it also outperforms second order DARTS while requiring less compute time, showcasing the benefit of geometry-aware optimization. Finally, EDARTS achieve comparable performance to the reported architecture for state-of-the-art method XNAS when evaluated using the stage 3 training routine of DARTS. Following XNAS, we also perform an extended evaluation of the best architecture found by EDARTS with AutoAugment, cosine power annealing , cross-entropy with label smoothing, and trains for 1500 epochs. We evaluated the XNAS architecture using our implementation for a direct comparison and also to serve as a reproducibility check. EDARTS achieved a test error of 2.18% in the extended evaluation compared to 2.15% for XNAS in our reproduced evaluation; note the published test error for XNAS is 1.81%. To meet a higher bar for reproducibility we report'broad reproducibility' by repeating the entire pipeline from stage 1 to stage 3 for two additional sets of seeds. Our in Table 2 (see Appendix) show that EDARTS has lower variance across experiments than random search with weight sharing . However, we do observe non-negligible variance in the performance of the architecture found by different random seed initializations of the shared-weights network, necessitating running multiple searches before selecting an architecture. A OPTIMIZATION This section contains proofs and generalizations of the non-convex optimization in Section 3. We first gather some necessary definitions and from convex analysis. Definition A.1. Let X be a convex subset of a finite-dimensional real vector space and f: X → R be everywhere sub-differentiable. 1. For α > 0, f is α-strongly-convex w.r.t. norm · if ∀ x, y ∈ X we have Definition A.2. Let X be a convex subset of a finite-dimensional real vector space. The Bregman divergence induced by a strictly convex distance-generating function (DGF) ω: X → R is D ω (x||y) = ω(x) − ω(y) − ∇ω(y), x − y ∀ x, y ∈ X By definition, the Bregman divergence satisfies the following properties: 3. If ω is β-strongly-smooth w.r.t. norm · then so is D ω (·||y) ∀ y ∈ X. Furthermore, D ω (x||y) ≤ β 2 x − y 2 ∀ x, y ∈ X. Lemma A.1 (Three-Points Lemma). (, Lemma 9.11) For any DGF ω: X → R and all x, y, z ∈ X we have ∇ω(y) − ∇ω(x), z − x = D ω (z||x) + D ω (x||y) − D ω (z||y) Definition A.3. Let ω: X → R be a 1-strongly-convex DGF. Then for constant λ > 0 and an everywhere sub-differentiable function f: X → R the proximal operator is defined over x ∈ X as prox λ (x) = arg min Note that the prox operator is well-defined whenever f is β-strongly-smooth for some β < λ. We will also use the following notation for the proximal gradient operator: Note that the prox grad operator is always well-defined. Theorem A.1. (, Theorem 9.12) For any λ > 0, 1-strongly-convex DGF ω: X → R, and x ∈ X let f: X → R be an everywhere sub-differentiable function s.t. λf (·) + D ω (·||x) is convex over X. Then for x + = prox λ (x) and all u ∈ X we have Lemma A.2. For any λ > 0, 1-strongly-convex DGF ω: X → R, and x ∈ X let f: X → R be an everywhere sub-differentiable function s.t. Proof. Applying Theorem A.1 followed by Lemma A.1 yields Corollary A.1. For any λ > 0, 1-strongly-convex DGF ω: X → R, x ∈ X, and everywhere sub-differentiable function f: X → R we have for Because we consider constrained non-convex optimization, we cannot measure convergence to a stationary point by the norm of the gradient. Instead, we analyze the convergence proximal-mappingbased stationarity measures. The most well-known measure is the norm of the projected gradient , which in the unconstrained Euclidean case reduces to the norm of the gradient and in the general case measure the distance between the current point and a cone satisfying first-order optimality conditions (, Proposition 4.1). Our convergence hold for a stronger measure that we call the projected stationarity and which is inspired by the Bregman stationarity measure of but using the prox grad operator instead of the prox operator. Definition A.4. Let ω: X → R be a 1-strongly-convex DGF and f: X → R be an everywhere sub-differentiable function. Then for any λ > 0 we define the following two quantities: The following properties follow: We can also consider the Bregman stationarity measure of directly. As this measure depends on the prox operator, which is not always defined, we first state the notion of non-convexity that consider. Definition A.5. An everywhere sub-differentiable function f: X → R is (γ, ω)-relatively-weaklyconvex ((γ, ω)-RWC) for γ > 0 and ω: X → R a 1-strongly-convex DGF if f (·) + γω(·) is convex over X. Note that all γ-strongly-smooth functions are (γ, ω)-RWC. Note that (γ, ω)-RWC is a generalization of γ-strong-smoothness w.r.t. the norm w.r.t. which ω is strongly-convex. Furthermore, for such functions we can always define the prox operator for λ > γ, allowing us to also define the Bregman gradient below. Similarly to before, bounding the Bregman stationarity measure yields a stronger notion of convergence than the squared norm of the Bregman gradient. For the relationship between the Bregman stationarity measure and first-order optimality conditions see Zhang and He (2018, Equation 2.11). Definition A.6. Let ω: X → R be a 1-strongly-convex DGF and f: X → R be a (γ, ω)-RWC everywhere sub-differentiable function for some γ > 0. Then for any λ > γ we define the following two quantities: Here we prove our main optimization . We begin with a descent lemma guaranteeing improvement of a non-convex function due to approximately optimizing a strongly convex surrogate. Lemma A.3. Let ω: X → R be a 1-strongly-convex DGF and f: X → R be everywhere subdifferentiable. For some x ∈ X and ρ > 0 definef x (·) = f (·) + ρD ω (·||x) and letx ∈ X be a point s.t. Ef x (x) − min u∈Xfx (u) ≤ ε. Then 1. If f is β-strongly-smooth, ρ > β, and λ = Proof. Generalizing an argument in Agarwal et al. (2019, Theorem A.2), for x + ∈ X we have by strong-convexity off x that If f is β-strongly-smooth set x + = prox∇ λ (x), so that by Corollary A.1 we have In the other case of f being (γ, ω)-RWC set x + = prox λ (x), so that by Lemma A.2 we have We now turn to formalizing our multi-block setting and assumptions. Setting A.1. For i = 1,..., b let X i be a convex subset of a real vector space with an associated DGF ω i: X i → R that is 1-strongly-convex w.r.t. some norm · (i) over X i. We have an everywhere sub-differentiable function f: X → R over the product space Our main will hold for the case when the following general assumption is satisfied. We will later show how this assumption can follow from strong smoothness or relative weak convexity and existing algorithmic . Assumption A.1. In Setting A.1, for any given ε > 0 and each i ∈ [b] there exists a constant ρ i > 0 and an algorithm A i: X → X that takes a point x ∈ X and returns a pointx ∈ X satisfyinĝ x −i = x −i and where the subscript i selects block i, the subscript −i selects all blocks other than block i, and E Ai denotes expectation w.r.t. the randomness of algorithm A i and any associated stochastic oracles. Algorithm 2: Generic successive convex approximation algorithm for reaching a stationary point of the non-convex function in Setting A.1. Input: Point x ∈ X in the product space of Setting A.1. Algorithms A 1,..., Our main relies on the following simple lemma guaranteeing non-convex convergence for a generic measure satisfying guaranteed expected descent: Lemma A.4. In Setting A.1, for some ε > 0 and each λi (x) be any measure s.t. for some λ i and some algorithm A i: X → X we have Then the output x of Algorithm 2 satisfies where F = f (x ) − arg min u∈X f (x) and the expectation is taken over the sampling at each iteration, the sampling of the output, and the randomness of the algorithms and any associated stochastic oracles. Proof. Define Ξ t = {(ξ s, A ξs)} t s=1 and note that x (t+1) = A ξt (x (t) ). We then have In the single-block setting, Lemmas A.3 and A.4 directly imply the following guarantee: Theorem A.2. In Setting A.1 and under Assumption A.1, let b = 1 and ρ satisfy one of the following: 1. f: X → R is β-strongly-smooth and ρ = 2β. 2. f: X → R is (γ, ω)-RWC for some DGF ω: X → R and ρ = 2γ. Then Algorithm 2 returns a point x ∈ X satisfying one of the following (respectively, w.r.t. the above settings) for Here the expectation is taken over the randomness of the algorithm and oracle. We can apply a known for the strongly convex case to recover the rate of for non-convex stochastic mirror descent, up to an additional depending on ω: Corollary A.2. In Setting A.1 for b = 1 and (γ, ω)-RWC f, suppose we have access to f through a stochastic gradient oracle g(x) = E∇f (x) such that E g 2 * ≤ G 2. Let A: X → X be an algorithm that for any x ∈ X runs the Epoch-GD method of with total number of steps N, initial epoch length T 1 = 4 and initial learning rate η 1 = 1 γ onf x (·) = f (·) + 2γD ω (·||x). Then with N T calls to the stochastic gradient oracle Algorithm 2 returns a point x ∈ X satisfying and L ω the Lipschitz constant of ω w.r.t. · over X. So an expected ε-stationary-point, as measured by Proof. Apply Theorem 5 of together with the fact thatf x is γ-stronglyconvex w.r.t. · and its stochastic gradient is bounded by For the multi-block case our hold only the projected stationarity measure: Theorem A.3. In Setting A.1 and under Assumption A.1 assume f (·, x −i) is β-strongly-smooth w.r.t. Here the expectation is taken over the randomness of the algorithm and oracle and the projected stationarity measure ∆ λ is defined w.r.t. the Bregman divergence of the DGF ω(where prox∇ To apply Lemma A.4 in the multiblock setting it suffices to show that the sum of the projected stationarity measures on each block is equal to the projected stationarity measure induced by the sum of the DGFs. For some λ > 0 and any i ∈ [b] we have that and so Thus applying Lemma A.2 with λ = 1 4ρ yields the . In the following corollary we recover the rate of for non-convex blockstochastic mirror descent, up to an additional term depending on ω i: Corollary A.3. In Setting A.1 for β-strongly-smooth f, suppose we have access to f through a stochastic gradient oracle g(x) = E∇f (x) such that E g i i. For i ∈ [b] let A i: X → X be an algorithm that for any x ∈ X runs the Epoch-GD method of with total number of steps N, initial epoch length T 1 = 4 and initial learning rate η 1 = 1 γ on surrogate functionf x (·) = f (·, x −i) + 2βD ωi (·||x i). Then with N T calls to the stochastic gradient oracle Algorithm 2 returns a point x ∈ X satisfying We can specialize this to the architecture search setting where we have a configuration search space contained in the product of simplices induced by having n decisions with c choices each together with a parameter space bounded in Euclidean norm. Corollary A.4. Under the assumptions of Corollary A.3, suppose b = 2 and we have the following two geometries: Suppose the stochastic gradient oracle of f has bounded ∞ -norm σ 1 over X 1 and σ 2 over X 2. Then Algorithm 2 will return an expected ε-stationary point of f under the projected stationarity measure in a number of stochastic oracle calls bounded by This section contains proofs of the generalization in Section 4. We first describe the setting for which we prove our general . Setting B.1. Let C be a set of possible architecture/configurations of finite size such that each c ∈ C is associated with a parameterized hypothesis class H c ={h (c) w: X → Y: w ∈ W} for input space Z = X × Y and fixed set of possible weights W. We will measure the performance of a hypothesis h (c) w on an input z = (x, y) ∈ Z using z (w, c) = (h Finally, we will consider solutions of optimization problems that depend on the training data and architecture. Specifically, for any configuration c ∈ C and finite subset S ⊂ Z let W c (S) ⊂ W be the set of global minima of some optimization problem induced by S and c and let the associated version space be H c (S) = {h w : X → Y is determined by a choice of architecture c ∈ C and a set of network weights w ∈ W and the loss : Y × Y → {0, 1} is the zero-one loss. In the simplest case W c (S) is the set of global minima of the ERM problem min We now state the main assumption we require. Assumption B.1. In Setting B.1 there exists a good architecture c * ∈ C, i.e. one satisfying (w *, c *) ∈ arg min W×C D (w, c) for some weights w * ∈ W, such that w.p. 1 − δ over the drawing of training set T ∼ D m T at least one of the minima of the optimization problem induced by c * and T has low excess risk, i.e. ∃ w ∈ W c * (T) s.t. for some error function ε c *. Clearly, we prefer error functions ε c * that are decreasing in the number of training samples m T and increasing at most poly-logarithmically in 1 δ. This assumption requires that if we knew the optimal configuration a priori, then the provided optimization problem will find a good set of weights for it. We will show how, under reasonable assumptions, Assumption B.1 can be formally shown to hold in Settings B.2 and B.3. Our general will be stated in terms of covering numbers of certain function classes. Definition B.1. Let H be a class of functions from X to Y. For any ε > 0 the associated L ∞ covering number N (H, ε) of H is the minimal positive integer k such that H can be covered by k balls of L ∞ -radius ε. The following is then a standard in statistical learning theory (see e.g. Lafferty et al. (2010, Theorem 7.82 where we use the loss notation from Setting B.1. Before stating our theorem, we define a final quantity, which measures the log covering number of the version spaces induced by the optimization procedure over a given training set. Definition B.2. In Setting B.1, for any sample S ⊂ X × Y define the version entropy to be Λ(H, ε, S) = log N c∈C H c (S), ε. Theorem B.2. In Setting B.1 let (ŵ,ĉ) ∈ W × C be obtained as a solution to the following optimization problem: arg min Then under Assumption B.1 we have w.p. 1 − 3δ that each term of which can be bounded as follows: 1. Sinceŵ ∈ Wĉ(T) for someĉ ∈ C the hypothesis space can be covered by the union of the coverings of H c (T) over c ∈ C, so by Theorem B.1 we have that w. 2. By optimality of the pair (ŵ,ĉ) and the fact that w ∈ W c * (T) we have 3. Hoeffding's inequality yields V (w, c We can then directly apply Theorem B.2 and the fact that the version entropy is bounded by log |C| because the minimizer over the training set is always unique to get the following: Corollary B.2. In Setting B.2 let (ŵ,ĉ) ∈ W × C be obtained as a solution to the following optimization problem: arg min In the special case of kernel selection we can apply generalization for learning with random features to show that we can compete with the optimal RKHS from among those associated with one of the configurations (, Theorem 1): Corollary B.3. In Setting B.2, suppose each configuration c ∈ C is associated with a random Fourier feature approximation of a continuous shift-invariant kernel that approximates an RKHS H c. Suppose is the squared loss so that (ŵ,ĉ) ∈ W × C is obtained as a solution to the following optimization problem: In the case of neural architecture search we are often solving (unregularized) ERM in the inner optimization problem. In this case we can make an assumption weaker than Assumption B.1, namely that the set of empirical risk minimizers contains a solution that, rather than having low excess risk, simply has low generalization error; then applying Hoeffding's inequality yields the following: Corollary B.4. In Setting B.1 let (ŵ,ĉ) ∈ W × C be obtained as a solution to the following optimization problem: arg min Suppose there exists c * ∈ C satisfying (w *, c *) ∈ arg min W×C D (w, c) for some weights w * ∈ W such that w.p. 1 − δ over the drawing of training set T ∼ D m T at least one of the minima of the optimization problem induced by c * and T has low generalization error, i.e. ∃ w ∈ arg min w ∈W T (w, c *) s.t. Solvers for Ridge regression and logistic regression were from scikit-learn . For CIFAR-10 we use the kernel configuration setting from but replacing the regularization parameter by the option to use the Laplace kernel instead of Gaussian. The regularization was fixed to λ = 1 2 The split is 40K/10K/10K. For IMDB we consider the following configuration choices: For stages 2 and 3, we train each architecture for 600 epochs with the same hyperparameter settings as DARTS. For completeness, we describe the convolutional neural network search space considered. The set of operations O considered at each node include: 3 × 3 separable convolution, 5 × 5 separable convolution, 3×3 dilated convolution, 5×5 dilated convolution, max pooling, average pooling, identity. We use the same search space to design a "normal" cell and a "reduction" cell; the normal cells have stride 1 operations that do not change the dimension of the input, while the reduction cells have stride 2 operations that half the length and width dimensions of the input. In the experiments, for both cell types, we set N = 6 with 2 input nodes and 4 intermediate nodes, after which the output of all intermediate nodes are concatenated to form the output of the cell. We use EDARTS to train a smaller shared-weights network in the search phase with 8 layers and 24 initial channels instead of the 16 used by DARTS. Additionally, to more closely mirror the architecture used for evaluation in stage 2 and 3, we use an auxiliary head with weight 0.4 and scheduled path dropout of 0.2. For the EDARTS architecture updates, we use a learning rate of 0.2 for the normal cell and 0.6 for the reduction cell. All other hyperparameters are the same as DARTS: 50 training epochs, batch size of 64, gradient clipping of 5 for network weights, SGD with momentum set to 0.9 and learning rate annealed from 0.025 to 0.001 with cosine annealing , and weight decay of 0.0003. We use the same evaluation scheme as DARTS when retraining architectures from scratch. The larger evaluation network has 20 layers and 36 initial channels and is trained for 600 epochs using SGD with momentum set to 0.9, a batch size of 96, and a learning rate of 0.025 annealed down to 0; the gradient clipping scheduled drop path rate and weight decay are identical to the search phase. We also use an auxiliary head with a weight of 0.4 and cutout . We investigate whether weight-sharing implicitly regularizes the hypothesis space by examining the 2 norms and distance from initialization of the shared-weights network relative to that observed when training the best EDARTS architecture from scratch. We use the same network depth and hyperparameters as those used for the shared-weights network to train the fixed architecture. Figure 4 shows the percent difference in the norms between the fixed architecture and the shared-weights network pruned to just the operations kept for the fixed architecture. From the chart, we can see that both the 2 distance from initialization and the 2 norm of the shared-weights is smaller than that of a fixed network are higher than that of the shared-weights network by over 40%, suggesting weight-sharing acts as a form of implicit regularization. Table 2 show that EDARTS has lower variance across experiments than random search with weight sharing . However, we do observe non-negligible variance in the performance of the architecture found by different random seed initializations of the shared-weights network, necessitating running multiple searches before selecting an architecture.

represent the tldr for finding the corresponding scientific article

An analysis of the learning and optimization structures of architecture search in neural networks and beyond.

Represent the short summary for finding the corresponding scientific article (long)

We introduce the NLC, a metric that is cheap to compute in the networks randomly initialized state and is highly predictive of generalization, at least in fully-connected networks.


Question: Represent the scientific article to find a short TLDR

Question: Deep latent variable models have seen recent success in many data domains. Lossless compression is an application of these models which, despite having the potential to be highly useful, has yet to be implemented in a practical manner. We present'`Bits Back with ANS' (BB-ANS), a scheme to perform lossless compression with latent variable models at a near optimal rate. We demonstrate this scheme by using it to compress the MNIST dataset with a variational auto-encoder model (VAE), achieving compression rates superior to standard methods with only a simple VAE. Given that the scheme is highly amenable to parallelization, we conclude that with a sufficiently high quality generative model this scheme could be used to achieve substantial improvements in compression rate with acceptable running time. We make our implementation available open source at https://github.com/bits-back/bits-back. The connections between information theory and machine learning have long been known to be deep, and indeed the two fields are so closely related that they have been described as'two sides of the same coin' BID18. One particularly elegant connection is the essential equivalence between probabilistic models of data and lossless compression methods. The source coding theorem BID22 can be thought of as the fundamental theorem describing this idea, and Huffman coding BID13, arithmetic coding BID28 and the more recently developed asymmetric numeral systems BID3 ) are actual algorithms for implementing lossless compression, given some kind of probabilistic model. The field of machine learning has experienced an explosion of activity in recent years, and we have seen a number of papers looking at applications of modern deep learning methods to lossy compression. BID7 discusses applications of a deep latent Gaussian model to compression, with an emphasis on lossy compression. BID0, BID23, BID1, and BID19 all implement lossy compression using (variational) auto-encoder style models, and BID24 train a model for lossy compression using a GAN-like objective. Applications to lossless compression have been less well covered in recent works. We seek to advance in this direction, and we focus on lossless compression using latent variable models. The lossless compression algorithms mentioned above do not naturally cater for latent variables. However there is a method, known as'bits back coding' BID27 BID10, first introduced as a thought experiment, but later implemented in BID5 and BID4, which can be used to extend those algorithms to cope with latent variables. Although bits back coding has been implemented in restricted cases by BID4, there is no known efficient implementation for modern neural net-based models or larger datasets. There is, in fact, a fundamental incompatibility between bits back and the arithmetic coding scheme with which it has previously been implemented. We resolve this issue, describing a scheme that instead implements bits back using asymmetric numeral systems. We term this new coding scheme'Bits Back with ANS' (BB-ANS). Our scheme improves on existing implementations of bits back coding in terms of compression rate and code complexity, allowing for efficient lossless compression of arbitrarily large datasets with deep latent variable models. We demonstrate the efficiency of BB-ANS by losslessly compressing the MNIST dataset with a variational auto-encoder (VAE), a deep latent variable model with continuous latent variables BID15 BID20. As far as we are aware, this is the first time bits back coding has been implemented with continuous latent variables. We find that BB-ANS with a VAE outperforms generic compression algorithms for both binarized and raw MNIST, even with a very simple model architecture. We extrapolate these to predict that the performance of BB-ANS with larger, state of the art models would be significantly better than generic compression algorithms. In this section we describe bits back coding, a method for lossless compression of data using a latent variable model. Before we describe bits back itself, we briefly discuss methods for encoding a stream of data given a fully observed model, a task sometimes referred to as'range coding' or'entropy coding'. We do not go into detail about the algorithms or their implementation, but describe the high level characteristics necessary for understanding bits back. For brevity, in the following sections we use simply log to refer to the base 2 logarithm, usually denoted log 2. Message lengths are measured in bits. Suppose that someone ('the sender') has a sequence of randomly distributed symbols, s = (s 1, ..., s N), with each s n drawn from a finite alphabet A n, which they would like to communicate to someone else ('the receiver') in as few bits as possible. Suppose that sender and receiver have access to a probabilistic model p for each symbol in the sequence given the previous, and can compute the mass p(s n = k | s 1, . . ., s n−1) for each k ∈ A n, n ∈ {1, . . ., N}.Arithmetic coding (AC) and asymmetric numeral systems (ANS) are algorithms which solve this problem, providing an encoding from the sequence s to a sequence of bits (referred to as the 'message'), and a decoding to recover the original data s. Both AC and ANS codes have message length equal to the'information content' h(s) − log p(s) of the sequence plus a small constant overhead of around 2 bits. By Shannon's Source Coding Theorem, the expected message length can be no shorter than the entropy of the sequence s, defined by H [s] E[h(s)], and thus AC and ANS are both close to optimal BID22 BID18. For long sequences the small constant overhead is amortized and has a negligible contribution to the compression rate. Critically for bits back coding, AC and ANS differ in the order in which messages are decoded. In AC the message is FIFO, or queue-like. That is, symbols are decoded in the same order to that in which they were encoded. ANS is LIFO, or stack-like. Symbols are decoded in the opposite order to that in which they were encoded. Note that the decoder in these algorithms can be thought of as a mapping from i.i.d. bits with p(b i = 0) = p(b i = 1) = 1 2 to a sample from the distribution p. Since we get to choose p, we can also think of ANS/AC as invertible samplers, mapping from random bits to samples via the decoder and back to the same random bits via the encoder. For a far more detailed introduction to arithmetic coding, see BID28, for asymmetric numeral systems, see BID3. We now give a short description of bits back coding, similar to those that have appeared in previous works. For a more involved derivation see Appendix A. We assume access to a coding scheme such as AC or ANS which can be used to encode and decode symbols according to any distribution. We will return to the question of which is the correct coding scheme to use in Section 2.4. Suppose now a sender wishes to communicate a symbol s 0 to a receiver, and that both sender and receiver have access to a generative model with a latent variable, y. For now we take y to be discrete, we address continuous latents in Section 2.5.1. Suppose both sender and receiver can compute the forward probabilities p(y) and p(s | y), and also have access to an approximate posterior q(y | s). Bits back coding allows the sender and receiver to efficiently encode and decode the symbol s 0.We must assume that, as well as the sample s 0, the sender has some extra bits to communicate. The sender can decode these extra bits to generate a sample y 0 ∼ q(y | s 0). Then they can encode the symbol s 0 according to p(s | y 0) and the latent sample according to p(y). The receiver then does the inverse to recover the latent sample and the symbol. The extra bits can also be recovered by the receiver by encoding the latent sample according to q(y | s 0). We can write down the expected increase in message length (over the extra bits): DISPLAYFORM0 This quantity is equal to the negative of the evidence lower bound (ELBO), sometimes referred to as the'free energy' of the model. A great deal of recent research has focused on inference and learning with approximate posteriors, using the ELBO as an objective function. Because of the above equivalence, methods which maximize the ELBO for a model are implicitly minimizing the message length achievable by bits back coding with that model. Thus we can draw on this plethora of existing methods when learning a model for use with bits back, safe in the knowledge that the objective function they are maximizing is the negative expected message length. If we wish to encode a sequence of data points, we can sample the extra bits for the first data point at random. Then we may use the encoded first data point as the extra information for the second data point, the encoded second data point as the extra information for the third, and so on. This daisychain-like scheme was first described by BID4, and was called'bits-back with feedback'. We refer to it simply as'chaining'.As BID4 notes, chaining cannot be implemented directly using AC, because of the order in which data must be decoded. Frey gets around this by implementing what amounts to a stack-like wrapper around AC, which incurs a cost both in code complexity and, importantly, in compression rate. The cost in compression rate is a of the fact that AC has to be'flushed' in between each iteration of bits back, and each flush incurs a cost which is implementation dependent but typically between 2 and 32 bits. The central insight of this work is to notice that the chaining described in the previous section can be implemented straightforwardly with ANS with zero compression rate overhead per iteration. This is because of the fact that ANS is stack-like by nature, which resolves the problems that occur if one tries to implement bits back chaining with AC, which is queue-like. We now describe this novel method, which we refer to as'Bits Back with ANS' (BB-ANS).We can visualize the stack-like state of an ANS coder as where the dashed line symbolizes the encoding/decoding end or'top' of the stack. When we encode a symbol s onto the stack we effectively add it to the end, ing in a'longer' state and when we decode (or equivalently, sample) a symbol t from the stack we remove it from the same end, ing in a'shorter' state, plus the symbol that we decoded. Table 1 shows the states of the sender as they encode a sample, using our bits back with ANS algorithm, starting with some'extra information' as well as the sample s 0 to be encoded. Table 1: Sender encodes a symbol s 0 using Bits Back with ANS. This process is clearly invertible, by reversing the order of operation and replacing encodes with decodes and sampling with encoding. Furthermore it can be repeated; the ANS stack at the end of encoding is still an ANS stack, and therefore can be readily used as the extra information for encoding the next symbol. The algorithm is compatible with any model whose prior, likelihood and (approximate) posterior can be encoded and decoded with ANS. A simple Python implementation of both the encoder and decoder of BB-ANS is given in Appendix C. A number of factors can affect the efficiency of compression with BB-ANS, and mean that in practice, the coding rate will never be exactly equal to the ELBO. For any algorithm based on AC/ANS, the fact that all probabilities have to be approximated at finite precision has some detrimental effect. When encoding a batch of only a small number of i.i.d. samples, with no'extra information' to communicate, the inefficiency of encoding the first datapoint may be significant. In the worst case, that of a batch with only one datapoint, the message length will be equal to the log joint, log p(s 0, y 0). Note that optimization of this is equivalent to maximum a posteriori (MAP) estimation. However, for a batch containing more than one image, this effect is amortized. FIG0 shows an example with 30 samples, where BB-ANS appears to perform well. Below we discuss two other issues which are specific to BB-ANS. We investigate the magnitude of these effects experimentally in Section 3.2. We find that when compressing the MNIST test set, they do not significantly affect the compression rate, which is typically close to the negative ELBO in our experiments. Bits back coding has previously been implemented only for models with discrete latent variables, in BID4. However, many successful latent variable models utilize continuous latents, including the VAE which we use in our experiments. We present here a derivation, based on BID18, of the surprising fact that continuous latents can be coded with bits back, up to arbitrary precision, without affecting the coding rate. We also briefly discuss our implementation, which as far as we are aware is the first implementation of bits back to support continuous latents. Further discussion can be found in Appendix B.We can crudely approximate a continuous probability distribution, with density function p, with a discrete distribution by partitioning the real line into'buckets' of equal width δy. Indexing the buckets with i ∈ I, we assign a probability mass to each bucket of P (i) ≈ p(y i)δy, where y i is some point in the i th bucket (say its centre).During bits back coding, we discretize both the prior and the approximate posterior using the same set of buckets. We use capital P and Q to denote discrete approximations. Sampling from the discrete approximation Q(i | s) uses approximately log(q(y i | s)δy) bits, and then encoding according to the discrete approximation to the prior P costs approximately log(p(y i)δy) bits. The expected message length for bits back with a discretized latent is therefore DISPLAYFORM0 The δy terms cancel, and thus the only cost to discretization from the discrepancy between our approximation and the true, continuous, distribution. However, if the density functions are sufficiently smooth (as they are in a VAE), then for small enough δy the effect of discretization will be negligible. Note that the number of bits required to generate the latent sample scales with the precision − log δy, meaning reasonably small precisions should be preferred in practice. Furthermore, the benefit from increasing latent precision past a certain point is negligible for most machine learning model implementations, since they operate at 32 bit precision. In our experiments we found that increases in performance were negligible past 16 bits per latent dimension. In our implementation, we divide the latent space into buckets which have equal mass under the prior (as opposed to equal width). This discretization is simple to implement and computationally efficient, and appears empirically to perform well. However, further work is required to establish whether it is optimal in terms of the trade-off between compression rate and computation. In our description of bits back coding in Section 2, we noted that the'extra information' needed to seed bits back should take the form of'random bits'. More precisely, we need the of mapping these bits through our decoder to produce a true sample from the distribution q(y | s). A sufficient condition for this is that the bits are i.i.d. Bernoulli distributed with probability 1 2 of being in each of the states 0 and 1. We refer to such bits as'clean'.During chaining, we effectively use each compressed data point as the seed for the next. Specifically, we use the bits at the top of the ANS stack, which are the of coding the previous latent y 0 according to the prior p(y). Will these bits be clean? The latent y 0 is originally generated as a sample from q(y | s 0). This distribution is clearly not equal to the prior, except in degenerate cases, so naively we wouldn't expect encoding y 0 according to the prior to produce clean bits. However, the true sampling distribution of y 0 is in fact the average of q(y | s 0) over the data distribution. That is, q(y) q(y | s)p(s)ds. This is referred to in BID11 as the'average encoding distribution'.If q is equal to the true posterior, then evidently q(y) ≡ p(y), however in general this is not the case. BID11 measure the discrepancy empirically using what they call the'marginal KL divergence' KL[q(z) p(z)], showing that this quantity contributes significantly to the ELBO for three different VAE like models learned on MNIST. This difference implies that the bits at the top the ANS stack after encoding a sample with BB-ANS will not be perfectly clean, which could adversely impact the coding rate. We demonstrate the BB-ANS coding scheme using a VAE. This model has a multidimensional latent with standard Gaussian prior and diagonal Gaussian approximate posterior: DISPLAYFORM0 We choose an output distribution (likelihood) p(s | y) suited to the domain of the data we are modelling (see below). The usual VAE training objective is the ELBO, which, as we noted in Section 2.2, is the negative of the expected message length with bits back coding. We can therefore train a VAE as usual and plug it into the BB-ANS framework. We consider the task of compressing the MNIST dataset BID17. We first train a VAE on the training set and then compress the test using BB-ANS with the trained VAE.The MNIST dataset has pixel values in the range of integers 0,..., 255. As well as compressing the raw MNIST data, we also present for stochastically binarized MNIST BID21. For both tasks we use VAEs with fully connected generative and recognition networks, with ReLU activations. For binarized MNIST the generative and recognition networks each have a single deterministic hidden layer of dimension 100, with a stochastic latent of dimension 40. The generative network outputs logits parameterizing a Bernoulli distribution on each pixel. For the full (non-binarized) MNIST dataset each network has one deterministic hidden layer of dimension 200 with a stochastic latent of dimension 50. The output distributions on pixels are modelled by a beta-binomial distribution, which is a two parameter discrete distribution. The generative network outputs the two beta-binomial parameters for each pixel. Instead of directly sampling the first latents at random, to simplify our implementation we instead initialize the BB-ANS chain with a supply of'clean' bits. We find that around 400 bits are required for this in our experiments. The precise number of bits required to start the chain depends on the entropy of the discretized approximate posterior (from which we are initially sampling).We report the achieved compression against a number of benchmarks in Table 2. Despite the relatively small network sizes and simple architectures we have used, the BB-ANS scheme outperforms benchmark compression schemes. While it is encouraging that even a relatively small latent variable model can outperform standard compression techniques when used with BB-ANS, the more Table 2: Compression rates on the binarized MNIST and full MNIST test sets, using BB-ANS and other benchmark compression schemes, measured in bits per dimension. We also give the negative ELBO value for each trained VAE on the test set.important observation to make from Table 2 is that the achieved compression rate is very close to the value of the negative test ELBO seen at the end of VAE training. In particular, the detrimental effects of finite precision, discretizing the latent (Section 2.5.1) and of less'clean' bits (Section 2.5.2) do not appear to be significant. Their effects can be seen in FIG2, accounting for the small discrepancy of around 1% between the negative ELBO and the achieved compression. Implementing a state-of-the-art latent variable model is not the focus of this work. However, as shown in our experiments, BB-ANS can compress data to sizes very close to the negative ELBO. This means that we can predict the best currently achievable compression using BB-ANS from the reported values of the negative ELBO for state-of-the-art latent variable models. We consider PixelVAE BID8, a latent variable model with close to state-of-the-art . We use their reported ELBO on binarized MNIST and the 64 × 64 ImageNet dataset introduced in van den BID26.The predictions are displayed in Table 3, and show that BB-ANS with PixelVAE may have a significantly better compression rate than existing schemes. These predictions are based on the assumption that the discrepancy between compression rate and ELBO will remain small for larger models. We believe this assumption is reasonable, since from the point of view of BB-ANS there are no fundamental differences, apart from dimensionality, between a complex, hierarchical VAE such as PixelVAE and the simple VAEs which we used for our experiments. We leave the experimental verification of these predictions to future work. Another potential extension of BB-ANS is to time series latent variable models such as hidden Markov models, or latent Gaussian state space models such as those studied in BID14. Such models could, in principal, be coded with BB-ANS, but the number of'extra bits' needed in a naive implementation scales with the length of the chain (the total time for a time series model), which could lead to a highly sub-optimal compression rate in practice. It would be useful to have a method for'interleaving' bits back with the time steps of the model, however it is unclear whether this is possible, and we leave deeper exploration of this problem to future work. Table 3: Predicted compression of BB-ANS with PixelVAE against other schemes, measured in bits per dimension. Modern machine learning models are optimized to exploit batch-parallelism and model-parallelism and run fastest on GPU hardware. Our current implementation of BB-ANS is written in pure Python, is not parallelized and executes entirely on CPU. During encoding/decoding the compression/decompression code is a computational bottleneck, running orders of magnitude slower than the computations of the model probabilities. However, we believe that almost all of the computation in the algorithm could be executed in parallel, on GPU hardware, potentially relieving this bottleneck. Firstly, our encoder requires computation of the CDF and inverse CDF of the distributions in the model. In the case of a VAE model of binarized MNIST, these are Gaussian and Bernoulli distributions. CDFs and inverse CDFs are already implemented to run on GPU, for many standard distributions, including Gaussian and Bernoulli, in various widely used machine learning toolboxes. Less trivial is the ANS algorithm. However, ANS is known to be amenable to parallelization. Techniques for parallel implementation are discussed in BID6, and BID16 presents an open source GPU implementation. We leave the performance optimization of BB-ANS, including adapting the algorithm to run on parallel architectures, to future work, but we are optimistic that the marriage of models which are optimized for parallel execution on large datasets with a parallelized and optimized BB-ANS implementation could yield an extremely high performance system. A neural net based model such as a VAE may have many thousands of parameters. Although not the focus of this work, the cost of communicating and storing a model's parameters may need to be considered when developing a system which uses BB-ANS with a large scale model. However, we can amortize the one-time cost of communicating the parameters over the size of the data we wish to compress. If a latent variable model could be trained such that it could model a wide class of images well, then BB-ANS could be used in conjunction with such a model to compress a large number of images. This would make the cost of communicating the model weights worthwhile to reap the subsequent gains in compression. Efforts to train latent variable models to be able to model such a wide range of images are currently of significant interest to the machine learning community, for example on expansive datasets such as ImageNet BID2 ). We therefore anticipate that this is the most fruitful direction for practical applications of BB-ANS.We also note that there have been many recent developments in methods to decrease the space required for neural network weights, without hampering performance. For example, methods involving quantizing the weights to low precision BID9 BID25, sometimes even down to single bit precision BID12, are promising avenues of research that could significantly reduce the cost of communicating and storing model weights. Probabilistic modelling of data is a highly active research area within machine learning. Given the progress within this area, it is of interest to study the application of probabilistic models to lossless compression. Indeed, if practical lossless compression schemes using these models can be developed then there is the possibility of significant improvement in compression rate over existing methods. We have shown the existence of a scheme, BB-ANS, which can be used for lossless compression using latent variable models. We demonstrated BB-ANS by compressing the MNIST dataset, achieving compression rates superior to generic algorithms. We have shown how to handle the issue of latent discretization. Crucially, we were able to compress to sizes very close to the negative ELBO for a large dataset. This is the first time this has been achieved with a latent variable model, and implies that state-of-the-art latent variable models could be used in conjunction with BB-ANS to achieve significantly better lossless compression rates than current methods. Given that all components of BB-ANS are readily parallelizable, we believe that BB-ANS can be implemented to run on GPU hardware, yielding a fast and powerful lossless compression system. We present here a more detailed derivation of bits back coding. As before, suppose that a sender and receiver wish to communicate a symbol s 0, and they both have access to a generative model with a latent variable, y. Suppose both sender and receiver can compute the forward probabilities p(y) and p(s | y). How might they communicate a sample s 0 from this model?Naively, the sender may draw a sample y 0 from p(y), and encode both y 0 and s 0 according to the forward model, p(y) and p(s | y 0) respectively. This would in a message length of − log p(y 0) + log p(s 0 | y 0) bits. The receiver could then decode according to the forward model by first decoding y 0 according to p(y) and then decoding s 0 according to p(s | y 0). However, they can do better, and decrease the encoded message length significantly. Firstly, if there is some other information which the sender would like to communicate to the receiver, then we may use this to our advantage. We assume the other information takes the form of some random bits. As long as there are sufficiently many bits, the sender can use them to generate a sample y 0 by decoding some of the bits to generate a sample from p(y), as described in Section 2.1. Generating this sample uses − log p(y 0) bits. The sender can then encode y 0 and s 0 with the forward model, and the message length will be − log p(y 0) + log p(s 0 | y 0) as before. But now the receiver is able to recover the other information, by first decoding s 0 and y 0, and then encoding y 0, reversing the decoding procedure from which the sample y 0 was generated, to get the'bits back'. This means that the net cost of communicating s 0, over the other information is DISPLAYFORM0 Secondly, note that we can choose any distribution for the sender to sample y 0 from, it does not have to be p(y), and it may vary as a function of s 0. If we generalize and let q(· | s 0) denote the distribution that we use, possibly depending functionally on s 0, we can write down the expected message length: DISPLAYFORM1 This quantity is equal to the negative of the evidence lower bound (ELBO), sometimes referred to as the'free energy' of the model. Having recognized this equivalence, it is straightforward to show using Gibbs' inequality that the optimal setting of q is the posterior p(y | s 0), and that with this setting the message length is DISPLAYFORM2 This is the information content of the sample s 0, which by the source coding theorem is the optimal message length. Thus bits back can achieve an optimal compression rate, if sender and receiver have access to the posterior. In the absence of such a posterior (as is usually the case), then an approximate posterior must be used. We note that BID1 and BID19 approach lossless compression with latent variables by generating a latent from an approximate posterior, and encoding according to the prior and likelihood as described above, but not recovering the bits back. BID1 mention that the cost of coding the hierarchical distribution is only a small fraction of the total coding cost in their setting. This small fraction upper bounds the potential gains from using bits back coding. However, their approach is sub-optimal, even if only slightly, and in the common case where more than one data-point is being encoded they would gain a better compression rate by using BB-ANS. As we discussed in Section 2.1, the coding scheme we wish to use, ANS, is defined for symbols in a finite alphabet. If we wish to encode a continuous variable we must restrict it to such a finite alphabet. This amounts to discretizing the continuous latent space. In choosing our discretization, it is important to note the following:• The discretization must be appropriate for the densities that will use it for coding. For example, imagine we were to discretize such that all but one of our buckets were in areas of very low density, with just one bucket covering the area with almost all of the density. This would in almost all of the latent variables being coded as the same symbol (corresponding to the one bucket with the majority of the density). Clearly this cannot be an efficient discretization.• The prior p(y) and the approximate posterior q(y | s) must share the same discretization.• The discretization must be known by the receiver before seeing data, since the first step of decoding is to decode y 0 according the prior. We propose to satisfy these considerations, by using the maximum entropy discretization of the prior, p(y), to code our latent variable. This amounts to allocating buckets of equal mass under the prior. We visualize this for a standard Gaussian prior in FIG3. Having the discretization be a function of the prior (which is fixed) allows the receiver to know the discretization up front, which we have noted is necessary. This would not be true for a discretization that depended on the posterior. This discretization is appropriate for coding according to the prior, since we are maximizing the entropy for this density. However, it is not obvious that it will be appropriate for coding according to the posterior, which it must also be used for. Note that we can write our the expected message length (negative ELBO) for a single data point as:L(q) = −E q (y | s 0) log p(s 0 | y) + KL[q(y | s 0) p(y)]We can see that minimizing this objective encourages the minimization of the KL divergence between the posterior and the prior. Therefore a trained model will generally have a posterior'close' (in a sense defined by the KL divergence) to the prior. This indicates that the maximum entropy discretization of the prior may also be appropriate for coding according to the posterior. C BB-ANS PYTHON IMPLEMENTATION FIG4 shows code implementing BB-ANS encoding (as described in Table 1) and decoding in Python. Since the message is stack-like, we use the Pythonic names'append' and'pop' for encoding and decoding respectively. Notice that each line in the decoding'pop' method precisely inverts an operation in the encoding'append' method. The functions to append and pop from the prior, likelihood and posterior could in principle use any LIFO encoding/decoding algorithm. They may, for example, do ANS coding according to a sophisticated autoregressive model, which would be necessary for coding using PixelVAE. The only strict requirement is that each pop function must precisely invert the corresponding append function. For more detail, including an example implementation with a variational auto-encoder model (VAE), see the repository https://github.com/bits-back/bits-back.

Represent this text to find a much much longer article about it with say hundreds-thousands of words

We do lossless compression of large image datasets using a VAE, beat existing compression algorithms.

It's important to consider optimization in function space, not just parameter space. We introduce a learning rule that reduces distance traveled in function space, just like SGD limits distance traveled in parameter space.