Question: Represent the text to find a very short summary (1-2 sentences max):

Question: We tackle the problem of modeling sequential visual phenomena. Given examples of a phenomena that can be divided into discrete time steps, we aim to take an input from any such time and realize this input at all other time steps in the sequence. Furthermore, we aim to do this \textit{without} ground-truth aligned sequences --- avoiding the difficulties needed for gathering aligned data. This generalizes the unpaired image-to-image problem from generating pairs to generating sequences. We extend cycle consistency to \textit{loop consistency} and alleviate difficulties associated with learning in the ing long chains of computation. We show competitive compared to existing image-to-image techniques when modeling several different data sets including the Earth's seasons and aging of human faces. Image-to-image translation has gained tremendous attention in recent years. A pioneering work by shows that it is possible to realize a real image from one domain as a highly realistic and semantically meaningful image in another when paired data between the domains are available. Furthermore, CycleGAN extended the image-to-image translation framework in an unpaired manner by relying on the ability to build a strong prior in each domain based off generative adversarial networks (GANs, ) and enforcing consistency on the cyclic transformation from and to a domain. Methods similar to CycleGAN have also been developed roughly around the same time. Since its birth, CycleGAN has become a widely adopted technique with applications even beyond computer vision . However, CycleGAN family models are still somewhat limited since they only handle the translation problem (directly) between two domains. Modeling more than two domains would require separate instantiations of CycleGAN between any two pairs of domains -ing in a quadratic model complexity. A major recent work, StarGAN , addresses this by facilitating a fully connected domain-translation graph, allowing transformation between two arbitrary domains with a single model. This flexibility, however, appears restricted to domains corresponding to specific attribute changes such as emotions and appearance. Within nature, a multitude of settings exist where neither a set of pairs nor a fully-connected graph are the most natural representations of how one might proceed from one domain to another. In particular, many natural processes are sequentialand therefore the translation process should reflect this. A common phenomena modeled as an image-to-image task is the visual change of natural scenes between two seasons, e.g., Winter and Summer. This neglects the fact that nature first proceeds to Spring after Winter and Fall after Summer and therefore the pairing induces a very discontinuous reflection of the underlying process. Instead, we hope that by modeling a higher resolution discretization of this process, the model can more realistically approach the true model while reducing the necessary complexity of the model. It is difficult to obtain paired data for many image-to-image problems. Aligned sequential are even more difficult to come by. Thus, it is more plausible to gather a large number of examples from each step (domain) in a sequence without correspondences between the content of the examples. Therefore, we consider a setting similar to unpaired image-to-image transformation where we only have access to unaligned examples from each time step of the sequence being modeled. Given an example from an arbitrary point in the sequence, we then generate an aligned sequence over all other time steps -expecting a faithful realization of the image at each step. The key condition that required is that after generating an entire loop (returning from the last domain to the input domain), one should expect to return to the original input. This is quite a weak condition and promotes model flexibility. We denote this extension to the cycle consistency of as loop consistency and therefore name our approach as Loop-Consistent Generative Adversarial Networks (LoopGAN). This is a departure from many image-to-image approaches that have very short (usually length 2) paths of computation defining what it means to have gone "there and back", e.g. the ability to enforce reconstruction or consistency. Since we do not have aligned sequences, the lengths of these paths for LoopGAN are as large as the number of domains being modeled and require different approaches to make learning feasible. These are not entirely different from the problems that often arise in recurrent neural networks and we can draw similarities to our model as a memory-less recurrent structure with applied to images. We apply our method to the sequential phenomena of human aging and the seasons of the Alps with extensive comparisons with baseline methods for image-to-image translation. We also present additional on gradually changing azimuth angle of chairs and gradual change of face attributes to showcased the flexibility of our model. We show favorable against baseline methods for image-to-image translation in spite of allowing for them to have substantially larger model complexity. This is consistent with how CycleGAN is trained where two cycles are included. Generative Adversarial Networks Generative adversarial networks (GANs, ) implicitly model a distribution through two components, a generator G that transforms a sample from a simple prior noise distribution into a sample from the learned distribution over observable data. An additional component known as the discrimintor D, usually a classifier, attempts to distinguish the generations of G with samples from the data distribution. This forms a minimax game from which both G and D adapt to one another until some equilibrium is reached. Unpaired Image-to-Image Transformation As an extension to the image-to-image translation framework (pix2pix,), proposed CycleGAN which has a similar architecture as in but is able to learn transformation between two domains without paired training data. To achieve this, CycleGAN simultaneously train two generators, one for each direction between the two domains. Besides the GAN loss enforced upon by domain-wise discriminators, the authors proposed to add a cycle-consistency loss which forces the two generators to be reversible. Similar to pix2pix, this model aims at learning a transformation between two domains and cannot be directly applied in multi-domain setting that involves more than two domains. Concurrent to CycleGAN, proposed a method named UNIT that implicitly achieves alignment between two domains using a VAE-like structure where both domains share a common latent space. Furthermore, StarGAN ( ) proposed an image-to-image translation model for multiple domains. A single network takes inputs defining the source image and desired domain transformation, however, it has been mainly shown to be successful for the domains consisting of facial attributes and expressions. The problem of learning non-deterministic multi-modal transformation between two image domains has made progress in recent years (. The common approach that achieves good performance is to embed the images for both domains into a shared latent space. At test time, an input image in the source domain is first embedded into the shared latent space and decoded into the target domain conditioned on a random noise vector. These models avoid one-to-one deterministic mapping problem and are able to learn different transformations given the same input image. However, these models are developed exclusively for two-domain transformation and cannot be directly applied to problems with more than two domains. Style Transfer A specific task in image-to-image transformation called style transfer is broadly defined as the task of transforming a photo into an artistic style while preserving its content . Common approaches use a pre-trained CNN as feature extractor and optimize the output image to match low-level features with that of style image and match high-level features with that of content image . A network architecture innovation made popular by this field known as AdaIn combines instance normalization with learned affine parameters. It needs just a small set of parameters compared to the main network weights achieve different style transfers within the same network. It also shows great potential in improving image quality for image generation and image-to-image transformation . Face Aging Generating a series of faces in different ages given a single face image has been widely studied in computer vision. State-of-the-art methods use a combination of pre-trained age estimator and GAN to learn to transform the given image to different ages that are both age-accurate and preserve original facial structure. They rely heavily on a domain-specific age estimator and thus have limited application to the more general sequential image generation tasks that we try to tackle here. Video Prediction Video prediction attempts to predict some number of future frames of a video based on a set of input frames . Full videos with annotated input frames and target frames are often required for training these models. A combination of RNN and CNN models has seen success in this task . Predictive vision techniques (; ;) that use CNN or RNN to generate future videos also require aligned video clips in training. A recent work added a GAN as an extra layer of supervision for learning human trajectories. At a high level, video prediction can be seen as a supervised setting of our unsupervised task. Moreover, video prediction mostly aims at predicting movement of objections rather than transformation of a still object or scene which is the focus of our task. We formulate our method and objectives. Consider a setting of n domains, X 1,..., X n where i < j implies that X i occurs temporally before X j. This defines a sequence of domains. To make this independent of the starting domain, we additionally expect that can translate from X n to X 1 -something a priori when the sequence represents a periodic phenomena. We define a single generator G(x, i) where i ∈ {1, . . ., n} and x ∈ X i. Then, a translation between two domains X i and X j of an input x i ∈ X i is given by repeated applications of G in the form of G j−i (x i, i) (allowing for incrementing the second argument modulo n + 1 after each application of G). By applying G to an input n times, we have formed a direct loop of translations where the source and target domains are equal. While we use a single generator, we make use of n discriminators where D i is tasked with discriminating between a translation from any source domain to X i. Since we are given only samples from each domain X i, we refer to each domain X i = {x j} Ni j=1 as consisting of N i examples from the domain X i with data distribution p data (x i). Suppose x i ∼ p data (x i). Then we expect that for all other domains j, G ||j−i|| (x i, i) should be indistinguishable under D j from (true) examples drawn from p data (x j). Additionally, each D j should aim to minimize the ability for G to generate examples that it cannot identify as fake. This forms the adversarial objective for a specific domain as: where G * denotes iteratively applying G until x j is transformed into domain X i, i.e. ||j − i|| times. Taking this over all possible source domains, we get an overall adversarial objective as: where q(i) is a prior on the set of domains, eg. uniform. Within, an adversarial loss was supplemented with a cycle consistency loss that ensured applying the generator from domain A to domain B followed by applying a separate generator from B to A acts like an identity function. However, LoopGAN only has a single generator and supports an arbitrary number of domains. Instead, we build a loop of computations by applying the generator G to a source image n times (equal to the number of domains being modeled). This constitutes loop consistency and allows us to reduce the set of possible transformations learned to those that adhere to the consistency condition. Loop consistency takes the form of an L 1 reconstruction objective for a domain X i as: The combined loss of LoopGAN over both adversarial and loop-consistency losses can be written as: where λ weighs the trade-off between adversarial and loop consistency losses. An example instantiation of our framework for one loop in a four-domain problem is shown in Figure 1. We adopt the network architecture for style transfer proposed in as our generator. This architecture has three main components: a down-sampling module Enc(x), a sequence of residual blocks T (h, i), and an up-sampling module Dec(h). The generator G therefore is the composition G(x, i) = Dec(T (Enc(x), i)) where the dependence of T on i only relates to the stepspecific AdaIN parameters while all other parameters are independent of i. Following the notations from (;, let c7-k denote a 7 × 7 ConvReLU layer with k filters and stride 1, dk denote a 3 × 3 Conv-ReLU layer with k filters and stride 2, Rk denote a residual block with two 3 × 3 Conv-AdaIn-ReLU layers with k filters each, uk denotes a 3 × 3 fractional-strided-Conv-LayerNorm-ReLU layer with k filters and stride 1 2 . The layer compositions of modules are down-sampling: c7-32, d64, d128; residual blocks: R128 × 6; upsampling: u128, u64, c7-3. We use the PatchGAN discriminator architecture as : c4-64, c4-128, c4-256, c4-1, where c4-k denotes a 4 × 4 Conv-InstanceNorm-LeakyRelu(0.2) layer with k filters and stride 2. Suppose we wish to translate some x i ∈ X i to another domain X j. A naive approach would formulate this as repeated application of G, |j − i| times. However, referencing our definition of G, we can unroll this to find that we must apply Enc and Dec j − i times throughout the computation. However, Enc and Dec are only responsible for bringing an observation into and out of the space of T. This is not only a waste of computation when we only require an output at X j, but it has serious implications for the ability of gradients to propagate through the computation. Therefore, we implement G(x i, i) as: a single application of Enc(x i), j − i applications of T (h), and a single application of Dec(h). T is applied recurrently and the entire generator is of the form: We show in our ablation studies that this re-formulation is critical to the learning process and the ing quality of the transformations learned. Additionally, T (h, i) is given a a set of separate, learnable normalization (AdaIN ) parameters that it selects based off of i with all other parameters of T being stationary across time steps. The overall architecture is shown in Figure 2. For all datasets, the loop-consistency loss coefficient λ is set to 10. We use Adam optimizer ( ) with the initial learning rate of 0.0002, β 1 = 0.5, and β 2 = 0.999. We train the face aging dataset and Alps seasons dataset for 50 epochs and 70 epochs respectively with initial learning rate and linearly decay learning rate to 0 for 10 epochs for both datasets. We apply LoopGAN to two very different sequential image generation tasks: face aging and chaging seasons of scenery pictures. Baselines are built with two bi-domain models, CycleGAN and UNIT and also a general-purpose multi-domain model StarGAN . We are interested in the sequential transformation capabilities of separately trained bi-domains compared to LoopGAN. Therefore, for each of the two bi-domains models, we train a separate model between every pair of sequential domains, i.e. X i and X i+1 and additionally train a model between every pair (not necessarily sequential) domains X i and X j (i = j). The first approach allows us to build a baseline for sequential generation by chaining the (separately learned) models in the necessary order. For instance, if we have four domains: A, B, C, D, then we can train four separate CycleGAN (or UNIT) models: G AB, G BC, G CD, G DA and correctly compose them to replicate the desired sequential transformation. Additionally, we can train direct versions e.g. G AC of CycleGAN (or UNIT) for a more complete comparison against LoopGAN. We refer to composed versions of separately trained models as Chained-CycleGAN and Chained-UNIT depending on the base translation method used. Since StarGAN ( ) inherently allows transformation between any two domains, we can apply this in a chained or direct manner without any additional models needing to be trained. We adopt the UTKFace dataset for modeling the face aging task. It consists of over 20,000 face-only images of different ages. We divide the dataset into four groups in order of increasing age according to the ground truth age given in the dataset as A consisting of ages from 10-20, B containing ages 25-35, C containing ages 40-50, and D containing ages 50-100. The number of images for each group are 1531, 5000, 2245, 4957, respectively, where a 95/5 train/test split is made. The of LoopGAN generation are shown in on the left side in Figure 3. LoopGAN shows advantage over baseline models in two aspects. The overall facial structure is preserved which we believe is due to the enforced loop consistency loss. Moreover, LoopGAN is able to make more apparent age changes compared to the rest of baseline models. In order to quantitatively compare the amount of age change between models, we obtain an age distribution of generated images by running a pre-trained age estimator DEX . The estimated age distributions of generated images (from input test images) are compared against those of the train images in Figure 4. The age distribution of LoopGAN generated images is closer to that of the train images across all four age groups when compared to the baseline modelssuggesting that it more faithfully learns the sequential age distribution changes of the training data. We use the collected scenery photos of Alps mountains of four seasons from . They are ordered into a sequence starting from Spring (A), to Summer (B), Fall (C), and Winter (D). They each have approximately 1700 images and are divided into 95/5% training and test set. We show the in Figure 3. Overall, LoopGAN is able to make drastic season change while maintaining the overall structure of the input scenery images. To further quantify the generation , we conducted a user study with Amazon Mechanical Turk (AMT) To showcase the universality of our model, we apply LoopGAN to two additional datasets in four different sequential transformation tasks: chairs with different azimuth angles, and gradual change of face attributes in degree of smiling, gender feature, and hair color. The chairs dataset comes with ground truth azimuth angle and is divided into four sets each containing chairs facing a distinct direction. To obtain linear manifolds for the face attributes, we train a binary classifier with 0/1 labels available for each attribute (Liu et al.) and use the predicted probability to determine the position of an image on an attribute manifold. The are shown in Figure 5. We experiment with several network architecture variations and investigate their effect on generation quality. First, attention mechanisms have proven to be useful in GAN image generation . We added attention mechanim in both space and time dimension ), however we found that the network struggles to generate high quality image after adding this type of attention mechanism. We also noticed that mentioned that for down-sampling, it is better to use no normalization to preserve information from input image, and for up-sampling it is better to use layer-normalization for faster training and higher quality. We applied these changes and found that they indeed help the network produce better . The under these variations are shown in Figure 6.a (first three rows). Moreover, we show the importance of the recurrent form of T (h) discussed in Section 4.2. We compare the choice to invoke Enc and Dec at each time step versus applying them once with some number of recurrent applications of T in Figure 6.a (last row) and show the poor quality observed when performing the loop naively. Lastly, we calculate the parameter count of generator networks compared in the face aging and season change experiments above and show that our final generator network architecture is parameterefficient compared to baseline models in Figure 6.b. For completeness, we also include a selection of failure cases in table in Figure 7. (viewed in color). In both cases, the highlighted generated images (the first column in (a) and the last column in (b)) bear some semantic dissimilarity to the input images. It seems that sometimes the network overfit to more drastic transformations that only preserve overall facial structure and orientation but neglects all other features. We proposed an extension to the family of image-to-image translation methods when the set of domains corresponds to a sequence of domains. We require that the translation task can be modeled as a consistent loop. This allows us to use a shared generator across all time steps leading to significant efficiency gains over a nave chaining of bi-domain image translation architectures. Despite this, our architecture shows favorable when compared with the classic CycleGAN family algorithms.

LoopGAN extends cycle length in CycleGAN to enable unaligned sequential transformation for more than two time steps.

We learn an unsupervised learning algorithm that produces useful representations from a set of supervised tasks. At test-time, we apply this algorithm to new tasks without any supervision and show performance comparable to a VAE.


Question: We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm to accelerate DNN training. Our algorithm uses large mini-batches to compute an approximate solution quickly and then refines it by averaging the weights of multiple models computed independently and in parallel. The ing models generalize equally well as those trained with small mini-batches but are produced in a substantially shorter time. We demonstrate the reduction in training time and the good generalization performance of the ing models on the computer vision datasets CIFAR10, CIFAR100, and ImageNet. Stochastic gradient descent (SGD) and its variants are the de-facto methods to train deep neural networks (DNNs). Each iteration of SGD computes an estimate of the objective's gradient by sampling a mini-batch of the available training data and computing the gradient of the loss restricted to the sampled data. A popular strategy to accelerate DNN training is to increase the mini-batch size together with the available computational resources. Larger mini-batches produce more precise gradient estimates; these allow for higher learning rates and achieve larger reductions of the training loss per iteration. In a distributed setting, multiple nodes can compute gradient estimates simultaneously on disjoint subsets of the mini-batch and produce a consensus estimate by averaging all estimates, with one synchronization event per iteration. Training with larger mini-batches requires fewer updates, thus fewer synchronization events, yielding good overall scaling behavior. Even though the training loss can be reduced more efficiently, there is a maximum batch size after which the ing model tends to have worse generalization performance (; ; ; ;). This phenomenon forces practitioners to use batch sizes below those that achieve the maximum throughput and limits the usefulness of large-batch training strategies. Stochastic Weight Averaging (SWA) ) is a method that produces models with good generalization performance by averaging the weights of a set of models sampled from the final stages of a training run. As long as the models all lie in a region where the population loss is mostly convex, the average model can behave well, and in practice, it does. We have observed that if instead of sampling multiple models from a sequence generated by SGD, we generate multiple independent SGD sequences and average models from each, the ing model achieves similar generalization performance. Furthermore, if all the independent sequences use small-batches, but start from a model trained with large-batches, the ing model achieves generalization performance comparable with a model trained solely with small-batches. Using these observations, we derive Stochastic Weight Averaging in Parallel (SWAP): A simple strategy to accelerate DNN training by better utilizing available compute resources. Our algorithm is simple to implement, fast and produces good with minor tuning. For several image classification tasks on popular computer vision datasets (CIFAR10, CIFAR100, and ImageNet), we show that SWAP achieves generalization performance comparable to models trained with small-batches but does so in time similar to that of a training run with large-batches. We use SWAP on some of the most efficient publicly available models to date, and show that it's able to substantially reduce their training times. Furthermore, we are able to beat the state of the art for CIFAR10 and train in 68% of the time of the winning entry of the DAWNBench competition. The mechanism by which the training batch size affects the generalization performance is still unknown. A popular explanation is that because of the reduced noise, a model trained using larger mini-batches is more likely to get stuck in a sharper global minima. In , the authors argue that sharp minima are sensitive to variations in the data because slight shifts in the location of the minimizer will in large increases in average loss value. However, if flatness is taken to be the curvature as measured by the second order approximation of the loss, then counterexamples exist. In , the authors transform a flat minimizer into a sharp one without changing the behavior of the model, and in , the authors show the reverse behavior when weight-decay is not used. In , the authors predict that the batch size can be increased up to a critical size without any drop in accuracy and empirically validate this claim. For example, the accuracy begins to drop for image classification on CIFAR10 when the batch sizes exceed 1k samples. They postulate that when the batch size is large, the mini-batch gradient is close to the full gradient, and further increasing the batch size will not significantly improve the signal to noise ratio. In , the authors argue that, for a fixed number of epochs, using a larger batch size implies fewer model updates. They argue that changing the number of updates impacts the distance the weights travel away from their initialization and that this distance determines the generalization performance. They show that by training with large-batches for longer times (thus increasing the number of updates), the generalization performance of the model is recovered. Even though this large-batch strategy generates models that generalize well, it does so in more time than the smallbatch alternative. Irrespective of the generalization performance, the batch size also affects the optimization process. In , the authors show that for convex functions in the over-parameterized setting, there is a critical batch size below which an iteration with a batch size of M is roughly equivalent to M iterations with a batch size of one, and batch-sizes larger than M do not improve the rate of convergence. Methods which use adaptive batch sizes exist (; ; ; ;). However, most of these methods are either designed for specific datasets or require extensive hyper-parameter tuning. Furthermore, they ineffectively use the computational resources by reducing the batch size during part of the training. Local SGD is a distributed optimization algorithm that trades off gradient precision with communication costs by allowing workers to independently update their models for a few steps before synchronizing. Post-local SGD is a variant, which refines the output of large-batch training with local-SGD. The authors have observed that the ing model has better generalization than the model trained with large-batches and that their scheme achieves significant speedups. In this manner Post-local SGD is of a very similar vein than the present work. However, while Post-local SGD lets the models diverge for T iterations where T is in the order of tens, SWAP averges the models once after multiple epochs. For example, in our Imagenet exeperiments (see Sec. 5) we average our models after tens of thousands of updates, while Post-local SGD does after at most 32. Because of this difference, we believe that the mechanisms that power the success of SWAP and Post-local SGD must be different and point to different phenomena in DNN optimization. Stochastic weight averaging (SWA) ) is a method where models are sampled from the later stages of an SGD training run. When the weights of these models are averaged, they in a model with much better generalization properties. This strategy is very effective and has been adopted in multiple domains: deep reinforcement learning (Nikishin et al.), semisupervised learning , Bayesian inference , lowprecision training. In this work, we adapt SWA to accelerate DNN training. We describe SWAP as an algorithm in three phases (see Algorithm 1): In the first phase, all workers train a single model by computing large mini-batch updates. Synchronization between workers is required at each iteration and a higher learning rate is used. In the second phase, each worker independently refines its copy of the model to produce a different set of weights. Workers use a smaller batch size, a lower learning rate, and different randomizations of the data. No synchronization between workers is required in this phase. The last phase consists of averaging the weights of the ing models and computing new batch-normalization statistics to produce the final output. Phase 1 is terminated before the training loss reaches zero or the training accuracy reaches 100% (for example, a few percentage points below 100%). We believe that stopping early precludes the optimization from getting stuck at a location where the gradients are too small and allows the following stage to improve the generalization performance. However, the optimal stopping accuracy is a hyper-parameter that requires tuning. During phase 2, the batch size is appropriately reduced and small-batch training is performed independently and simultaneously. Here, each worker (or a subset of them) performs training using all the data, but sampling in different random order. Thus, after the end of the training process, each worker (or subset) will have produced a different model. Figure 1 plots the accuracies and learning-rate schedules for a run of SWAP. During the large-batch phase (phase 1), all workers share a common model and have the same generalization performance. During the small-batch phase (phase 2) the learning rates for all the workers are the same but their testing accuracies differ as the stochasticity causes the models to diverge from each other. We also plot the test-accuracy of the averaged model that would were we to stop phase 2 at that point. Note that the averaged model performs consistently better than each individual model. Small-batch Phase Figure 1: Learning rate schedules and CIFAR10 test accuracies for workers participating in SWAP. The largebatch phase with synchronized models is followed by the small-batch phase with diverging independent models. The test accuracy of the averaged weight model is computed by averaging the independent models and computing the test loss for the ing model. To visualize the mechanism behind SWAP, we plot the error achieved by our test network on a plane that contains the outputs of the three different phases of the algorithm. Inspired by and, we pick orthogonal vectors u, v that span the plane which contains θ 1, θ 2, θ 3. We plot the loss value generated by model θ = θ 1 +αu+βv at the location (α, β). To plot a loss value, we first generate a weight vector θ, compute the batch-norm statistics for that model (through one pass over the training data), and then evaluate the test and train accuracies. In Figure 2, we plot the training and testing error for the CIFAR10 dataset. Here'LB' marks the output of phase one,'SGD' the output of a single worker after phase two, and'SWAP' the final Figure 2a, we observe that the level-sets of the training error (restricted to this plane) form an almost convex basin and that both the output of phase 1 ('LB') 2 and the output of one of the workers of phase 2 ('SGD') lie in the outer edges of the basin. Importantly, during phase 2 the model traversed to a different side of the basin (and not to the center). Also, the final model ('SWAP') is closer to the center of the basin. When we visualize these three points on the test loss landscape (Figure 2b), we observe that the variations in the topology of the basin cause the'LB' and'SGD' points to fall in regions of higher error. But, since the'SWAP' point is closer to the center of the basin, it is less affected by the change in topology. In Figure 3, we neglect the'LB' point and plot the plane spanned by three workers'SGD1','SGD2','SGD3'. In Figure 3a, we can observe that these points lie at different sides of the training error basin while'SWAP' is closer to the center. In Figure 3b, we observe that the change in topology causes the worker points to lie in regions of higher testing errors than'SWAP', which is again close to the center of both basins. For reference, we have also plotted the best model that can be generated by this region of the plane. In , the authors argue that in the later stages of SGD the weight iterates behave similar to an Ornstein Uhlenbeck process. So, by maintaining a constant learning rate the SGD iterates should reach a stationary distribution that is similar to a high-dimensional Gaussian. This distribution is centered at the local minimum, has a covariance that grows proportionally with the learning rate, inversely proportional to the batch size and has a shape that depends on both the Hessian of the mean loss and covariance of the gradient. The authors of argue that by virtue of being a high dimensional Gaussian all the mass of the distribution is concentrated near the'shell' of the ellipsoid, and therefore, it is unlikely for SGD to access the interior. They further argue that by sampling weights from an SGD run (leaving enough time steps between them) will choose weights that are spread out on the surface of this ellipsoid and their average will be closer to the center. Without any further assumptions, we can justify sampling from different SGD runs (as done in phase 2 during SWAP). As long as all runs start in the same basin of attraction, and provided the model from holds, all runs will converge to the same stationary distribution, and each run can generate independent samples from it. To win some intuition on the advantage that SWA and SWAP have over SGD, we measure the cosine similarity between the gradient descent direction, −g i, and the direction towards the output of SWAP, ∆θ = θ swap − θ i. In Figure 4, we see that the cosine similarity, ∆θ,−gi gi ∆θ, decreases as the training enters its later stages. We believe that towards the end of training, the angle between the gradient direction and the directions toward the center of the basin is large, therefore the process moves mostly orthogonally to the basin, and progress slows. However, averaging samples from different sides of the basin can (and does) make faster progress towards the center. In this section we evaluate the performance of SWAP for image classification tasks on the CIFAR10, CIFAR100, and ImageNet datasets. For the experiments in this subsection, we found the best hyper-parameters using grid searches (see Appendix A for details). We train using mini-batch SGD with Nesterov momentum (set to 0.9) and weight decay of 5×10 −4. We augment the data using cutout and use a fastto-train custom ResNet 9 from a submission 3 to the DAWNBench leaderboard (Coleman et al.). All experiments were run on one machine with 8 NVIDIA Tesla V100 GPUs and use Horovod to distribute the computation. All statistics were collected over 10 different runs. For these experiments, we used the following settings-SWAP phase one: 4096 samples per batch using 8 GPUs (512 samples per GPU). Phase one is terminated when the training accuracy reaches 98% (on average 108 epochs). SWAP phase two: 8 workers with one GPU each and 512 samples per batch for 30 epochs. The experiment that uses only large-batches had 4096 samples per batch across 8 GPUs and is run for 150 epochs. The experiments that use only small-batches had 512 samples per batch on 2 GPUs and is trained for 100 epochs. Table 1 compares the best test accuracies and corresponding training times for models trained with small-batch only, with large-batch only, and with SWAP. We report the average accuracy of the workers before averaging and the accuracy of the final model. Test Table 2 compares the best test accuracies and corresponding training times for models trained with only small-batches (for 150 epochs), with only large-batches (for 150 epochs), and with SWAP. For SWAP, we report test accuracies obtained using the last SGD iterate before averaging, and test accuracy of the final model obtained after averaging. We observe significant improvement in test accuracies after averaging the models. For both CIFAR 10 and CIFAR100, training with small-batches achieves higher testing accuracy than training with large-batches but takes much longer to train. SWAP, however, terminates in time comparable to the large-batch run but achieves accuracies on par (or better) than small batch training. Achieving state of the art training speeds for CIFAR10: At the time of writing the front-runner of the DAWNBench competition takes 37 seconds with 4 Tesla V100 GPUs to train CIFAR10 to 94% test accuracy. Using SWAP with 8 Tesla V100 GPUs, a phase one batch size of 2048 samples and 28 epochs, and a phase two batch size of 256 samples for one epoch is able to reach the same accuracy in 27 seconds. We use SWAP to accelerate a publicly available fast-to-train ImageNet model with published learning rate and batch size schedules 4. The default settings for this code modify the learning-rates and batch sizes throughout the optimization (see Figure 5). Our small-batch experiments train ImageNet for 28 epochs using the published schedules with no modification and are run on 8 Tesla V100 GPUs. Our large-batch experiments modify the schedules by doubling the batch size and doubling the learning rates (see Figure 5) and are run on 16 Tesla V100 GPUs. For SWAP phase 1, we use the large-batch settings for 22 epochs, and for SWAP phase 2, we run two independent workers each with 8 GPUs using the settings for small-batches for 6 epochs. We observe that doubling the batch size reduces the Top1 and Top5 test accuracies with respect to the small-batch run. SWAP, however, recovers the generalization performance at substantially reduced training times. Our are compiled in Table 3 (the statistics were collected over 3 runs). We believe it's worthy of mention that these accelerations were achieved with no tuning other than increasing the learning rates proportionally to the increase in batch size and reverting to the original schedule when transitioning between phases. taken from an existing DAWNBench submission. For a larger batch experiment, we double the batch size, double the number of GPUs and double the learning rate of the original schedule. For SWAP, we switch from the modified schedule to the original schedule as we move from phase 1 to phase 2. Top1 Accuracy (%) Top5 Accuracy (%) Training Time (min) SGD (small-batch) 76.14 ± 0.07 93. We now compare SWAP with SWA: the sequential weight averaging algorithm from. For the experiments in this section, we use the CIFAR100 dataset. We sample the same number of models for both SWA and SWAP and maintain the same number of epochs per sample. For SWA, we sample each model with 10 epochs in-between and average them to get the final model. For SWAP, we run 8 independent workers for 10 epochs each and use their average as the final model. We explore if SWA can recover the test accuracy of small-batch training on a large-batch training run. We use the same (large) batch size throughout. We follow an initial training cycle with cyclic learning rates (with cycles of 10 epochs) to sample 8 models (one from the end of each cycle). See Figure 6a for an illustration of the learning rate schedule. As expected we observe that the large-batch training run achieves lower training accuracy, but surprisingly SWA was unable to improve it (see Table 4, row 1). Large-batch followed by small-batch SWA: We evaluate the effect of executing SWA using smallbatches after a large-batch training run. We interrupt the large-batch phase at the same accuracy we interrupt phase 1 of our CIFAR100 experiment (Table 2). In this case, the small-batch phase uses a single worker and samples the models sequentially. SWA is able to reach the test accuracy of a small-batch run but requires more than three times longer than SWAP to compute the model (see Table 4, row 2). An illustration of the learning rate schedule is provided in Figure 6b. Small-batch SWA and SWAP: We start the SWA cyclic learning rate schedule from the best model found by solely small-batch training (table 2, row 1). Since the cycle length and cycle count are fixed, the only free parameter is the peak learning rate. We select this using a grid-search. Once the SWA schedule is specified, we re-use the peak learning rate settings in SWAP. We start phase two from the model that was generated as the output of phase 1 for the experiment on section 5.1 reported on table 2 rows 3 and 4. With these settings, small-batch SWA achieves better accuracy than SWAP (by around ∼ 0.9%) at 6.8x more training time. Next, we wish to explore the speed-up that SWAP achieves over SWA if the precision of SWA is set as a target. To that end, we relax the constraints on SWAP. By increasing the phase two schedule from one 10 epoch cycle to two 20 epoch cycles and sampling two models from each worker (16 models) the ing model achieved a test accuracy of 79.11% in 241 seconds or 3.5x less time. Test accuracy before averaging (%) We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm that uses a variant of Stochastic Weight Averaging (SWA) to improve the generalization performance of a model trained with large mini-batches. Our algorithm uses large mini-batches to compute an approximate solution quickly and then refines it by averaging the weights of multiple models trained using small-batches. The final model obtained after averaging has good generalization performance and is trained in a shorter time. We believe that this variant and this application of SWA are novel. We observed that using large-batches in the initial stages of training does not preclude the models from achieving good generalization performance. That is, by refining the output of a large-batch run, with models sampled sequentially as in SWA or in parallel as in SWAP, the ing model is able to perform as well as the models trained using small-batches only. We confirm this in the image classification datasets CIFAR10, CIFAR100, and ImageNet. Through visualizations, we complement the existing evidence that averaged weights are closer to the center of a training loss basin than the models produced by stochastic gradient descent. It's interesting to note that the basin into which the large mini-batch run is converging to seems to be the same basin where the refined models are found. So, it is possible that regions with bad and good generalization performance are connected through regions of low training loss and, more so, that both belong to an almost convex basin. Our method requires the choice of (at least) one more hyperparameter: the transition point between the large-batch and small-batch. For our experiments, we chose this by using a grid search. A principled method to choose the transition point will be the focus of future work. In future work we intend to explore the behavior of SWAP when used with other optimization schemes, such as Layer-wise Adaptive Rate Scaling (LARS) , mixed-precision training , post-local SGD or NovoGrad . The design of SWAP allows us to substitute any of these for the large-batch stage, for example, we can use local SGD to accelerate the first stage of SWAP by reducing the communication overhead. We provide the parameters used in the experiments of Section 5.1. These were obtained by doing independent grid searches for each experiment. For all CIFAR experiments, the momentum and weight decay constants were kept at 0.9 and 5 × 10 −4 respectively. Tables 5 and 6 list the remaining hyperparameters. When a stopping accuracy of 100% is listed, we mean that the maximum number of epochs were used.

We propose SWAP, a distributed algorithm for large-batch training of neural networks.

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

We propose a hypothesis for why gradient descent generalizes based on how per-example gradients interact with each other.


Question: Represent the text to find a very short summary (1-2 sentences max)

Question: A common way to speed up training of large convolutional networks is to add computational units. Training is then performed using data-parallel synchronous Stochastic Gradient Descent (SGD) with a mini-batch divided between computational units. With an increase in the number of nodes, the batch size grows. However, training with a large batch often in lower model accuracy. We argue that the current recipe for large batch training (linear learning rate scaling with warm-up) is not general enough and training may diverge. To overcome these optimization difficulties, we propose a new training algorithm based on Layer-wise Adaptive Rate Scaling (LARS). Using LARS, we scaled AlexNet and ResNet-50 to a batch size of 16K. Training of large Convolutional Neural Networks (CNN) takes a lot of time. The brute-force way to speed up CNN training is to add more computational power (e.g. more GPU nodes) and train network using data-parallel Stochastic Gradient Descent, where each worker receives some chunk of global mini-batch (see e.g. BID10 or BID4). The size of a chunk should be large enough to utilize the computational resources of the worker. So scaling up the number of workers in the increase of batch size. But using large batch may negatively impact the model accuracy, as was observed in BID10, BID14, BID8, BID6.Increasing the global batch while keeping the same number of epochs means that you have fewer iterations to update weights. The straight-forward way to compensate for a smaller number of iterations is to do larger steps by increasing the learning rate (LR). For example, BID10 suggests to linearly scale up LR with batch size. However using a larger LR makes optimization more difficult, and networks may diverge especially during the initial phase. To overcome this difficulty, BID4 suggested a "learning rate warm-up": training starts with a small LR, which is slowly increased to the target "base" LR. With a LR warm-up and a linear scaling rule, BID4 successfully trained ResNet-50 BID5 ] with batch B=8K, see also BID1 ]. Linear scaling of LR with a warm-up is the "state-of-the art" recipe for large batch training. We tried to apply this linear scaling and warm-up scheme to train AlexNet BID11 ] on ImageNet BID3 ], but scaling stopped after B=2K since training diverged for large LR-s. For B=4K the accuracy dropped from the baseline 57.6% (B=512) to 53.1%, and for B=8K the accuracy decreased to 44.8%. To enable training with a large LR, we replaced Local Response Normalization layers in AlexNet with Batch Normalization (BN) BID7 ]. We will refer to this models AlexNet-BN. BN improved model convergence for large LRs as well as accuracy: for B=8K the accuracy gap decreased from 14% to 2.2%.To analyze the training stability with large LRs we measured the ratio between the norm of the layer weights and norm of gradients update. We observed that if this ratio is too high, the training becomes unstable. On other hand, if the ratio is too small, then weights don't change fast enough. The layer with largest ||∇W || ||W || defines the global limit on the learning rate. Since this ratio varies a lot between different layers, we can speed-up training by using a separate LR for each layer. Thus we propose a novel Layer-wise Adaptive Rate Scaling (LARS) algorithm. There are two notable differences between LARS and other adaptive algorithms such as ADAM BID9 ) or RMSProp BID16 ): first, LARS uses a separate learning rate for each layer and not for each weight, which leads to better stability. And second, the magnitude of the update is defined with respect to the weight norm for better control of training speed. With LARS we trained AlexNet-BN and ResNet-50 with B=16K without accuracy loss. The training of CNN is done using Stochastic Gradient (SG) based methods. At each step t a minibatch of B samples x i is selected from the training set. The gradients of loss function ∇L(x i, w) are computed for this subset, and networks weights w are updated based on this stochastic gradient: DISPLAYFORM0 The computation of SG can be done in parallel by N units, where each unit processes a chunk of the mini-batch with B N samples. Increasing the mini-batch permits scaling to more nodes without reducing the workload on each unit. However, it was observed that training with a large batch is difficult. To maintain the network accuracy, it is necessary to carefully adjust training hyper-parameters (learning rate, momentum etc). BID10 suggested the following rules for training with large batches: when you increase the batch B by k times, you should also increase LR by k times while keeping other hyper-parameters (momentum, weight decay, etc) unchanged. The logic behind linear LR scaling is straight-forward: if you increase B by k times while keeping the number of epochs unchanged, you will do k times fewer steps. So it seems natural to increase the step size by k times. For example, let's take k = 2. The weight updates for batch size B after 2 iterations would be: DISPLAYFORM1 The weight update for the batch B 2 = 2 * B with learning rate λ 2: DISPLAYFORM2 will be similar if you take λ 2 = 2 * λ, assuming that ∇L(x j, w t+1) ≈ ∇L(x j, w t).Using the "linear LR scaling" BID10 trained AlexNet with batch B=1K with minor (≈ 1%) accuracy loss. The scaling of AlexNet above 2K is difficult, since the training diverges for larger LRs. It was observed that linear scaling works much better for networks with Batch Normalization (e.g. BID2). For example BID0 trained the Inception model with batch B=6400, and trained ResNet-152 for B=5K.The main obstacle for scaling up batch is the instability of training with high LR. BID6 tried to use less aggressive "square root scaling" of LR with special form of Batch Normalization ("Ghost Batch Normalization") to train AlexNet with B=8K, but still the accuracy (53.93%) was much worse than baseline 58%. To overcome the instability during initial phase, BID4 proposed to use LR warm-up: training starts with small LR, and then LR is gradually increased to the target. After the warm-up period (usually a few epochs), you switch to the regular LR policy ("multi-steps", polynomial decay etc). Using LR warm-up and linear scaling BID4 trained ResNet-50 with batch B=8K without loss in accuracy. These recipes constitute the current state-of-the-art for large batch training, and we used them as the starting point of our experiments. Another problem related to large batch training is so called "generalization gap", observed by BID8. They came to that "the lack of generalization ability is due to the fact that large-batch methods tend to converge to sharp minimizers of the training function." They tried a few methods to improve the generalization with data augmentation and warm-starting with small batch, but they did not find a working solution. We used BVLC 1 AlexNet with batch B=512 as baseline. Model was trained using SGD with momentum 0.9 with initial LR=0.02 and the polynomial (power=2) decay LR policy for 100 epochs. The baseline accuracy is 58.8% (averaged over last 5 epochs). Next we tried to train AlexNet with B=4K by using larger LR. In our experiments we changed the base LR from 0.01 to 0.08, but training diverged with LR > 0.06 even with warm-up 2. The best accuracy for B=4K is 53.1%, achieved for LR=0.05. For B=8K we couldn't scale-up LR either, and the best accuracy is 44.8%, achieved for LR=0.03 (see TAB0 (a) ).To stabilize the initial training phase we replaced Local Response Normalization layers with Batch Normalization (BN). We will refer to this model as AlexNet-BN.3. AlexNet-BN model was trained using SGD with momentum=0.9, weight decay=0.0005 for 128 epochs. We used polynomial (power 2) decay LR policy with base LR=0.02. The baseline accuracy for B=512 is 60.2%. With BN we could use large LR-s even without warm-up. For B=4K the best accuracy 58.9% was achieved for LR=0.18, and for B=8K the best accuracy 58% was achieved for LR=0.3. We also observed that BN significantly widens the range of LRs with good accuracy. Still there is a 2.2% accuracy loss for B=8K. To check if it is related to the "generalization gap" (Keskar et al. FORMULA0), we looked at the loss gap between training and testing (see FIG0). We did not find the significant difference in the loss gap between B=512 and B=8K. We conclude that in this case the accuracy loss was mostly caused by the slow training and was not related to a generalization gap. The standard SGD uses the same LR λ for all layers: w t+1 = w t − λ∇L(w t). When λ is large, the update ||λ * ∇L(w t)|| can become larger than ||w||, and this can cause the divergence. This makes the initial phase of training highly sensitive to the weight initialization and to initial LR. We found that the ratio of the L2-norm of weights and gradients ||w||/||∇L(w t)|| varies significantly between weights and biases, and between different layers. For example, let's take AlexNet after one iteration TAB1, "*.w" means layer weights, and "*.b" -biases). The ratio ||w||/||∇L(w)|| for the 1st convolutional layer ("conv1.w") is 5.76, and for the last fully connected layer ("fc6.w") -1345. The ratio is high during the initial phase, and it is rapidly decreasing after few epochs (see Figure 2). If LR is large comparing to the ratio for some layer, then training may becomes unstable. The LR "warm-up" attempts to overcome this difficulty by starting from small LR, which can be safely used for all layers, and then slowly increasing it until weights will grow up enough to use larger LRs. We would like to use different approach. We want to make sure that weights update is small comparing to the norm of weights to stabilize training DISPLAYFORM0 where η < 1 control the magnitude of update with respect to weights. The coefficient η defines how much we "trust" that the value of stochastic gradient ∇L(w l t) is close to true gradient. The η depends on the batch size. "Trust" η is monotonically increasing with batch size: for example for Alexnet for batch B = 1K the optimal η = 0.0002, for batch B = 4K -η = 0.005, and for B = 8K -η = 0.008. We implemented this idea through defining local LR λ l for each layer l: DISPLAYFORM1 where γ defines a global LR policy (e.g. steps, or exponential decay), and local LR λ l is defined for each layer through "trust" coefficient η < 1 4: DISPLAYFORM2 Note that now the magnitude of the update for each layer doesn't depend on the magnitude of the gradient anymore, so it helps to partially eliminate vanishing and exploding gradient problems. The network training for SGD with LARS are summarized in the Algorithm 1 5.LARS was designed to solve the optimization difficulties, and it does not replace standard regularization methods (weight decay, batch norm, or data augmentation). But we found that with LARS we can use larger weight decay, since LARS automatically controls the norm of layer weights: DISPLAYFORM3 where B is min-batch size, N -number of training epochs, and S -number of samples in the training set. Here we assumed that global rate policy starts from 1 and decrease during training over training interval [0, N * S/B]. BID4 and BID1 we used the second setup with an extended augmentation with variable image scale and aspect ratio similar to ]. The baseline top-1 accuracy for this setup is 75.4%. The accuracy with B=16K is 0.7-1.4% less than baseline. This gap is related to smaller number of steps. We will show in the next section that one can recover the accuracy by training for more epochs. When batch becomes large (32K), even models trained with LARS and large LR don't reach the baseline accuracy. One way to recover the lost accuracy is to train longer (see BID6). Note that when batch becomes large, the number of iteration decrease. So one way to try to improve the accuracy, would be train longer. For example for Alexnet and Alexnet-BN with B=16K, when we double the number of iterations from 7800 (100 epochs) to 15600 (200 epochs) the accuracy improved by 2-3% (see TAB4). The same effect we observed for Resnet-50: training for additional 100 epochs recovered the top-1 accuracy to the baseline 75.5%. In general we found that we have to increase the training duration to keep the accuracy. Consider for example Googlenet ]. As a baseline we trained BVLC googlenet 8 with batch=256 for 100 epoch. The top-1 accuracy of this model is 69.2%. Googlenet is deep, so in original paper authors used auxiliary losses to accelerate SGD. We used LARS to solve optimization difficulties so we don't need these auxiliary losses. The original model also has no Batch Normalization, so we used data augmentation for better regularization. The baseline accuracy for B=256 is 70.3% with extended augmentation and LARS. We found that Googlenet is very difficult to train with large batch even with LARS: we needed both large number of epoch and longer ramp-up to scale learning rate up (see TAB5). Large batch is a key for scaling up training of convolutional networks. The existing approach for large-batch training, based on using large learning rates, leads to divergence, especially during the initial phase, even with learning rate warm-up. To solve these difficulties we proposed the new optimization algorithm, which adapts the learning rate for each layer (LARS) proportional to the ratio between the norm of weights and norm of gradients. With LARS the magnitude of the update for each layer doesn't depend on the magnitude of the gradient anymore, so it helps with vanishing and exploding gradients. But even with LARS and warm-up we couldn't increase LR farther for very large batches, and to keep the accuracy we have to increase the number of epochs and use extensive data augmentation to prevent over-fitting.

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

A new large batch training algorithm  based on Layer-wise Adaptive Rate Scaling (LARS); using LARS, we scaled AlexNet  and ResNet-50 to a batch of 16K.

We propose a extension of the batch normalization, show a first-of-its-kind convergence analysis for this extension and show in numerical experiments that it has better performance than the original batch normalizatin.


Question: Represent the text to find a very short summary (1-2 sentences max)

Question: Finding an embedding space for a linear approximation of a nonlinear dynamical system enables efficient system identification and control synthesis. The Koopman operator theory lays the foundation for identifying the nonlinear-to-linear coordinate transformations with data-driven methods. Recently, researchers have proposed to use deep neural networks as a more expressive class of basis functions for calculating the Koopman operators. These approaches, however, assume a fixed dimensional state space; they are therefore not applicable to scenarios with a variable number of objects. In this paper, we propose to learn compositional Koopman operators, using graph neural networks to encode the state into object-centric embeddings and using a block-wise linear transition matrix to regularize the shared structure across objects. The learned dynamics can quickly adapt to new environments of unknown physical parameters and produce control signals to achieve a specified goal. Our experiments on manipulating ropes and controlling soft robots show that the proposed method has better efficiency and generalization ability than existing baselines.

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

Learning compositional Koopman operators for efficient system identification and model-based control.

We propose a stochastic differentiable forward dynamics predictor that is able to sample multiple physically plausible trajectories under the same initial input state and show that it can be used to train model-free policies more efficiently.


Question: Represent the article for finding a single-sentence TLDR

Question: We derive reverse-mode (or adjoint) automatic differentiation for solutions of stochastic differential equations (SDEs), allowing time-efficient and constant-memory computation of pathwise gradients, a continuous-time analogue of the reparameterization trick. Specifically, we construct a backward SDE whose solution is the gradient and provide conditions under which numerical solutions converge. We also combine our stochastic adjoint approach with a stochastic variational inference scheme for continuous-time SDE models, allowing us to learn distributions over functions using stochastic gradient descent. Our latent SDE model achieves competitive performance compared to existing approaches on time series modeling. Deterministic dynamical systems can often be modeled by ordinary differential equations (ODEs). For training, a memory-efficient implementation of the adjoint sensitivity method effectively computes gradients through ODE solutions with constant memory cost. Stochastic differential equations (SDEs) are a generalization of ODEs which incorporate instantaneous noise into their dynamics (; Øksendal, 2003). They are a natural fit for modeling phenomena governed by small and unobserved interactions. In this paper, we generalize the adjoint method to dynamics defined by SDEs ing in an approach which we call the stochastic adjoint sensitivity method. Building on theoretical advances by , we derive a memory-efficient adjoint method whereby we simultaneously reconstruct the original trajectory and evaluate the gradients by solving a backward SDE (in the sense of) whose formulation we detail in Section 3. Computationally, in order to retrace the original trajectory during the backward pass, we need to reuse noise samples generated in the forward pass. In Section 4, we give an algorithm that allows arbitrarily-precise querying of a Brownian motion realization at any time point, while only storing a single random seed. Overall, this in a constant-memory algorithm that approximates the gradient arbitrarily well as step size reduces by computing vector-Jacobian products a constant number of times per-iteration. See Table 2 for a comparison of our method against previous approaches in terms of asymptotic time and memory complexity. We incorporate SDEs into a stochastic variational inference framework, whereby we efficiently compute likelihood ratios and backpropagate through the evidence lower bound using our adjoint approach. This effectively generalizes existing model families such as latent ODEs and deep Kalman filters . We review works on learning ODEs and SDEs. We refer the reader to Appendix B on for stochastic flows and (backward) Stratonovich integrals. The adjoint sensitivity method is an efficient approach to solve optimization problems by considering the dual form . recently applied this idea to obtain gradients with respect to parameters of a neural network defining an ODE. The method is scalable due to its memory-efficiency, as intermediate computations need not be cached as in regular backpropagation . Recent works have considered SDEs whose drift and diffusion functions are defined by neural networks (a,b; ;). Consider a filtered probability space (Ω, F, {F t} t∈T, P ) on which an m-dimensional adapted Wiener process {W t} t∈T is defined. An Itô SDE defines a stochastic process {Z t} t∈T by where z 0 ∈ R d is a deterministic starting value, and b: are the drift and diffusion functions, respectively. Here, the second integral on the right hand side of is the Itô stochastic integral (Øksendal, 2003). When the coefficients are globally Lipschitz in both the state and time components, there exists a unique strong solution to the SDE (Øksendal, 2003). Therefore, one can consider coefficients defined by neural networks that have smooth activation functions (e.g. tanh) of the form b(z, t, θ) and σ (z, t, θ). This in a model known as the neural SDE. We derive a backward Stratonovich SDE for what we call the stochastic adjoint process. A direct implication of this is a gradient computation algorithm that works by solving a set of dynamics in reverse time and relies on vector-Jacobian products without storing intermediate computation. Recall from Appendix B.3, Φ s,t (z):= Z s,z t is the solution at time t when the process is started at z at time s, and its inverse is defined asΨ s,t (z):= Φ −1 s,t (z). Consider A s,t (z) = ∇(L(Φ s,t (z))), where L is a scalar loss function. The chain rule gives A s,t (z) = ∇L(Φ s,t (z))∇Φ s,t (z). LetÃ s,t (z):= A s,t (Ψ s,t (z)) = ∇L(z)∇Φ s,t (Ψ s,t (z)) = ∇L(z)K s,t (z). Note that A s,t (z) =Ã s,t (Φ s,t (z)). Since ∇L(z) is constant, we see that (Ã s,t (z),Ψ s,t (z)) satisfies the following backward SDE system by Lemma C.1 (cf. Appendix C) Since can be viewed as a single SDE (with smooth coefficients) for an augmented state,Ã s,T (z) also has a unique strong solution. Therefore, for t = 0, we may writẽ where W · = {W t} 0≤t≤T denotes the path of the Brownian motion and F: is a deterministic measurable function (the Itô map) (, Definition 10.9). The next theorem follows immediately from and the definition of F. Theorem 3.1: The theorem is a consequence of T ) and. This implies we may solve the dynamics starting from the end state of the forward solve Z 0,z T to obtain the gradient of the loss with respect to the starting value z. To obtain the gradient with respect to the parameters, we augment the original state with parameters. Algorithm 1 summarizes this assuming access to a black-box solver SDESolve. See details in Appendix C. Input: parameters θ, start time t 0, stop time t 1, final state z t 1, loss gradient ∂L/z t 1. Input: drift b(z, t, θ), diffusion σ(z, t, θ), Wiener process sample w(t). We present a data structure that allows arbitrarilyprecise query of the sample path of the Wiener process given a global random seed based on the Brownian tree construction. The data structure facilitates the adjoint method such that we can ensure the noise sample in the backward solve is the same as in the forward solve using a split pseudorandom random number generator (PRNG). We present the procedure in Algorithm D.2 and details in Appendix D. Consider the SDEs where are Lipschitz in both arguments. Suppose and define the prior and posterior processes, respectively. Additionally, assume there is a function u: for all x ∈ R d and t ∈ R. Then, the variational free energy can be written as where the expectation is taken over the distribution of the posterior process defined by, and y 1,..., y N are observations at times t 1,..., t N, respectively. To compute the gradient with respect to parameters, we need only augment the forward equation with an extra variable whose drift function returns 1 2 |u(Z t, t)| 2 and diffusion function is 0. In this case, the backward adjoint dynamics can be derived analogously using. Appendix E includes details. We verify our theory by comparing the gradients obtained by our stochastic adjoint framework against analytically derived gradients for chosen test problems with closed-form solutions. We then fit latent SDE models with our framework on two synthetic datasets and a real dataset, verifying that the variational inference framework promotes learning a generative model of time series. Due to space constraint, we refer the reader to Appendix F for on numerical studies and Appendix N for on synthetic data. We present only the on the motion capture dataset here. We experiment on a dataset extracted from the CMU motion capture library. We use the dataset adopted by which consists of 23 walking sequences of subject number 35 that is partitioned into 16 training, 3 validation, and 4 test sequences. We include the settings in Appendix O and report the test MSE here following Yıldız et al.. Appendix B. Additional Background t. For the diffusion functions σ 1,..., σ m, we will also write σ: R d → R d×m as the matrix-valued function obtained by stacking the component functions σ i in a columnwise fashion, and index its jth row and ith column by σ j,i. Among recent work on neural SDEs, none has enabled an efficient training framework. In particular, Tzen and Raginsky (2019a); considered computing the gradient by simulating the forward dynamics of an explicit Jacobian matrix the size either the squared number of parameters or the number of parameters times the number of states, building on the pathwise approach . By contrast, the approach we present only requires evaluating vector-Jacobian products a constant number of times with respect to the number of parameters and states, which has the same asymptotic time cost as evaluating the drift and diffusion functions, and can be done automatically by modern machine learning libraries (; ; ;). Our stochastic adjoint sensitivity method involves stochastic processes running forward and backward in time. The Stratonovich stochastic integral, due to its symmetry, gives nice expressions for the backward dynamics and so is more convenient for our purpose. Our can be applied straightforwardly to Itô SDEs as well using a conversion (see e.g. (, Sec. 2 Following the treatment of Kunita , we introduce the forward and backward Stratonovich integrals. Let {F s,t} s≤t;s,t∈T be a two-sided filtration, where F s,t is the σ-algebra generated by For a continuous semimartingale {Y t} t∈T adapted to the forward filtration {F 0,t} t∈T, the Stratonovich stochastic integral is defined as where denotes the size of largest segment of the partition, and the limit is to be interpreted in the L 2 sense. The Itô integral uses instead the left endpoint Y t k rather than the average. In general, the Itô and Stratonovich integrals differ by a term of finite variation. To define the backward Stratonovich integral, we consider the backward Wiener process {W t} t∈T defined asW t = W t − W T for all t ∈ T that is adapted to the backward filtration {F t,T} t∈T. For a continuous semimartingaleY t adapted to the backward filtration, we define where Π = {0 = t N < · · · < t 0 = T} is a partition, and the limit is again in the L 2 sense. It is well known that an ODE defines a flow of diffeomorphisms . Here we consider the stochastic analogue for the Stratonovich SDE Throughout the paper, we assume b, σ are of class C ∞,1 b, so that the SDE has a unique strong solution. Let Φ s,t (z):= Z s,z t be the solution at time t when the process is started at z at time s. Given a realization of the Wiener process, this defines a collection S = {Φ s,t} s≤t;s,t∈T of continuous maps from R d to itself. The following theorem shows that these maps are diffeomorphisms and that they satisfy backward SDEs. Theorem B.1 (Thm. 3.7.1 ): (i) With probability 1, the collection S = {Φ s,t} s≤t;s,t∈T satisfies the flow property Moreover, each Φ s,t is a smooth diffeomorphism from R d to itself. We thus call S the stochastic flow of diffeomorphisms generated by the SDE. (ii) The backward flowΨ s,t:= Φ −1 s,t satisfies the backward SDE: for all z ∈ R d and s, t ∈ T such that s ≤ t a.s. Note that the coefficients in and differ by only a negative sign. This symmetry is due to our use of the Stratonovich integral (see Figure 2). Figure 2: Negating the drift and diffusion functions for an Itô SDE and simulating backwards from the end state gives the wrong solution. Negating the drift and diffusion functions for the converted Stratonovich SDE, however, gives the correct path when simulated in reverse time. We present our main contribution, i.e. the stochastic analog of the adjoint sensitivity method for SDEs. We use to derive another backward Stratonovich SDE for what we call the stochastic adjoint process. The direct implication of this is a gradient computation algorithm that works by solving a set of dynamics in reverse time and relies on vector-Jacobian products without storing intermediate computation produced in the forward pass. The goal is to derive the stochastic adjoint process {∂L/∂Z t} t∈T that can be simulated by evaluating only vector-Jacobian products, where L = L(Z T) is a scalar loss of the terminal state Z T. The main theoretical is Theorem 3.1. We first derive a backward SDE for the process {∂Z T /∂Z t} t∈T, assuming that Z t =Ψ t,T (Z T) for a deterministic Z T ∈ R d that does not depend on the realized Wiener process. We then extend to the case where In the latter case, the ing value cannot be interpreted as the solution to a backward SDE anymore due to loss of adaptiveness; instead we will formulate the using the Itô map . Finally, we extend the state of Z to include parameters and obtain the gradient with respect to them. We first derive the SDE for the Jacobian matrix of the backward flow. Consider the stochastic flow generated by the backward SDE as in Theorem B.1(ii). Let J s (z):= ∇Ψ s,T (z), then it satisfies the backward SDE for all s ≤ t and x ∈ R d a.s. Furthermore, let K s,t (z) = J s,t (z) −1, we have for all s ≤ t and x ∈ R d a.s. The proof included in Appendix I relies on Itô's lemma in the Stratonovich form (, Theorem 2.4.1). This lemma considers only the case where the endpoint z is fixed and deterministic. Now we compose the state process (represented by the flow) and the loss function L. WritingX s = (Ã s,T (z),Ψ s,T (z) ) as the augmented process, the system is a backward Stratonovich SDE of the form. As a has a unique strong solution. Without loss of generality, assume t = 0. Since admits a strong solution, we may writeÃ where W · = {W t} 0≤t≤T denotes the path of the Brownian motion and is a deterministic measurable function (the Itô map) (, Definition 10.9). Intuitively, F can be thought as an algorithm that computes the solution to the backward SDE given the position z at time T and the realized Brownian path. Similarly, we let G be the solution map for the forward flow. Immediately, we arrive at Theorem 3.1. In practice, we compute solutions to SDEs with numerical solvers F h and G h, where h = T /N denotes the mesh size of a fixed grid 1. The approximate algorithm thus outputs The following theorem provides sufficient conditions for convergence. Suppose the schemes F h and G h satisfy the following conditions: converge to 0 in probability as h → 0, and (ii) for any M > 0, we have sup |z|≤M |F h (z, W ·) − F(z, W ·)| → 0 in probability as h → 0. Then, for any starting point z of the forward flow, we have in probability as h → 0. For details and the proof see Appendix J. Usual schemes such as the Euler-Maruyama and Milstein method satisfy condition (i). Indeed, they converge pathwise (i.e. almost surely) with explicit rates for any fixed starting point . While condition (ii) is rather strong, we note that the SDEs considered here have smooth coefficients and thus the solutions enjoy nice regularity properties in the starting position. Therefore, it is reasonable to expect that the corresponding numerical schemes to also behave nicely as a function of both the mesh size and the starting position. To the best of our knowledge this property is not considered at all in the literature on numerical methods for SDEs (where the initial position is fixed), but is crucial in the proof of Theorem C.2. Detailed analysis for specific schemes is beyond the scope of this paper and is left for future research. So far we have derived the gradient of the loss with respect to the initial state. We can extend these to give gradients with respect to parameters of the drift and diffusion functions by treating them as an additional part of the state whose dynamics has zero drift and diffusion. We summarize this in Algorithm 1 2, assuming access to a numerical solver SDESolve. Note for the Euler-Maruyama scheme, the most costly terms to compute a t ∂b/∂θ and a t ∂σ i /∂θ can be evaluated by calling vjp(a t, b, θ) and vjp(a t, σ i, θ), respectively. In principle, we can simulate the forward and backward adjoint dynamics with any high-order solver of choice. However, in practice, to obtain a strong numerical solution 3 with order beyond 1/2, we need to simulate multiple integrals of the Wiener process such as These random variables are difficult to simulate exactly and costly to approximate using truncated infinite series . Note that even though the backward SDE for the stochastic adjoint does not have diagonal noise, it satisfies a commutativity property (Rößler, 2004) when the SDE of the original 1. We may also use adaptive solvers . 2. We use row vector notation here. 3. A numerical scheme is of strong order p if E [|XT − XNη|] ≤ Cη p for all T > 0, where Xt and XNη are respectively the coupled true solution and numerical solution, N and η are respectively the iteration index and step size such that N η = T, and C is independent of η. dynamics has diagonal noise. In this case, we can safely adopt certain numerical schemes of strong order 1.0 (e.g. Milstein and stochastic Runge-Kutta (Rößler, 2010)) without approximating multiple integrals or the Lévy area during simulation. We verify this formally in Appendix K. We have implemented several SDE solvers in PyTorch which include Euler-Maruyama, Milstein, and stochastic Runge-Kutta schemes with adaptive time-stepping using a PI controller . In addition, following torchdiffeq , we have created a user-friendly subclass of torch.autograd.Function that facilitates gradient computation using our stochastic adjoint framework when the neural SDE is implemented as a subclass of torch.nn.Module. We include a short code snippet covering the main idea of the stochastic adjoint in Appendix L and plan to release all code after the double-blind reviewing process. The formulation of the adjoint ensures it can be numerically integrated by merely evaluating dynamics cheaply defined by vector-Jacobian products, as opposed to whole Jacobians. However, the backward-in-time nature also introduces the additional difficulty that the same Wiener process sample path in the forward pass has to be queried again during the backward pass. Naïvely storing Brownian motion increments and related quantities (e.g. Lévy area approximations) not only implies a large memory consumption but also disables using adaptive time-stepping numerical integrators, where the evaluation timestamps in the backward pass may be different from those in the forward pass. To overcome this issue, we combine Brownian trees with splittable Pseudorandom number generators (PRNGs) and obtain a data structure that allows querying values of the Wiener process path at arbitrary times with logarithmic time cost with respect to some error tolerance. Lévy's Brownian bridge states that given a start time t s and end time t e along with their respective Wiener process values w s and w e, the marginal of the process at time t ∈ (t s, t e) is a normal distribution: We can recursively apply this formula to evaluate the process at the midpoint of any two distinct timestamps where the values are already known. Constructing the whole sample path of a Wiener process in this manner in what is known as the Brownian tree . We assume access to a splittable PRNG (Claessen and Pałka, 2013), which has an operation split that deterministically generates two (or more) keys 4 using an existing key. In addition, we assume access to an operation BrownianBridge which samples from given a key. To obtain the Wiener process value at a specific time, the seeded Brownian tree works by recursively sampling according to the Brownian tree with keys split from those of parent nodes, assuming the values at some initial and terminal times are known. The algorithm terminates when the current time under consideration is within a certain error tolerance of the desired time. We outline the full procedure in Algorithm D.2. This algorithm has constant memory cost. For fixed-step-size solvers, the tolerance that the tree will be queried at will scale as 1/L, where L is the number of steps in the solver. Thus the complexity per-step will scale as log L. Note that the variational free energy can be derived from Girsanov's change of measure theorem . To efficiently Monte Carlo estimate this quantity and its gradient, we simplify the equation by noting that for a one-dimensional process {V t} t∈T adapted to the filtration generated by a one-dimensional Wiener process {W t} t∈T, if Novikov's condition (Øksendal, 2003) is satisfied, then the process defined by the Itô integral t 0 V s dW s is a Martingale (Øksendal, 2003). Hence, E T 0 u(Z t, t) dW t = 0, and To Monte Carlo simulate the quantity in the forward pass along with the original dynamics, we need only extend the original augmented state with an extra variable L t such that the new drift and diffusion functions for the new augmented state By, the backward SDEs of the adjoint processes become In this case, neither do we need to actually simulate the backward SDE of the extra variable nor do we need to simulate its adjoint. Moreover, when considered as a single system for the augmented adjoint state, the diffusion function of the backward SDE satisfies the commutativity property. We consider three carefully designed test problems (examples 1-3 ; details in Appendix M) all of which have closed-form solutions. We compare the gradient computed from simulating our stochastic adjoint process using the Milstein scheme against the gradient evaluated by analytically solving the equations. Figure F (a) shows that for test example 1, the error between the adjoint gradient and analytical gradient decreases as the fixed step size decreases. One phenomenon not covered by our theory is that the error can be indeed be controlled by the adaptive solver. This is shown by the fact that for all three test problems, the mean-square error across dimensions tends to be smaller as the absolute tolerance is reduced (see Figure F (c, f, j) ). However, we note that the Number of Function Evaluations (NFEs) tends to be much larger than that in the ODE case , which is expected given the inherent roughness of Brownian motion paths. Sensitivity Analysis for SDEs. Gradient computation is closely related to sensitivity analysis. Computing gradients with respect to parameters of vector fields of an SDE has been extensively studied in the stochastic control literature . In particular, for low dimensional problems, this is done effectively using dynamic programming and finite differences . However, both approaches scale poorly with the dimensionality of the parameter vector. Analogous to REINFORCE (or score-function estimator) (; ;), as for some random variable H 5. However, H usually depends on the density of Z T with respect to the Lebesgue measure which can be difficulty to compute. extended this approach by weakening a non-degeneracy condition using Mallianvin calculus. Closely related to the current submission is the pathwise method , which is the continuous-time analog of the reparameterization trick . Existing methods in this regime (a; ;) all require simulating a forward SDE where each step requires computing entire Jacobian matrices. This computational cost is prohibitive for high-dimensional systems with a large number of parameters. Based on the Euler discretization, considered storing the intermediate values and performing reverse-mode automatic differentiation. They named this method the adjoint approach, which, by modern standards, is a form of "backpropagation 5. The random variable H is not unique. through the operations of a numerical solver". We comment that this approach, despite widely adopted in the field of finance for calibrating market models , has high memory cost, and relies on a fixed step size Euler-Maruyama discretization. This approach was used by to parameterize the drift and diffusion of an SDE using Gaussian processes. Backward SDEs. Our backward SDE for the stochastic adjoint process relies on the notion of backward SDEs by which is based on two-sided filtrations. This is different from the more traditional notion of backward SDEs where only a single filtration is defined . Based on the latter notion, forward-backward SDEs (FBSDEs) have been proposed to solve the stochastic optimal control problem . However, simulating FBSDEs is costly due to the need to estimate conditional expectations in the backward pass. Estimating conditional expectations, however, is a direct consequence of the appearance of an auxiliary process from the Martingale representation theorem . For notational convenience we suppress z and W ·. Bounding I 1. Let > 0 be given. Since G h → G in probability, there exist M 1 > 0 and h 0 > 0 such that Since the SDE defines a stochastic flow of diffeomorphisms, there exists a finite random variable C 2 such that sup |z|≤2M 1 |∇ z F| ≤ C 2, and there exists M 2 > 0 such that P(|C 2 | > M 2) <. Given M 2, there exists h 1 > 0 such that Now suppose h ≤ min{h 0, h 1}. Then, by the union bound, with probability at least 1 − 4, we have On this event, we have Thus, we have shown that I 1 converges to 0 in probability as h → 0. Bounding I 2. The idea is similar. By condition (ii), we have in probability. Using this and condition (i), for given > 0, there exist M > 0 and h 0 > 0 such that for h ≥ h 0, we have with probability at least 1 −. On this event, we have Thus I 2 also converges to 0 in probability. Recall the Stratonovich SDE with drift and diffusion functions b, σ 1,..., σ m ∈ R d × R → R d governed by a set of parameters θ ∈ R p. Consider the augmented state composed of the original state and parameters Y t = (Z t, θ). The augmented state satisfies a Stratonovich SDE with the drift function f (y, t) = (b(z, t), 0 p ) and diffusion functions. By and, the dynamics for the adjoint process of the augmented state is characterized by the backward SDE: By definitions of f and g i, the Jacobian matrices ∇f (x, s) and ∇g i (x, s) can be written as: Thus, we can write out the backward SDEs for the adjoint processes of the state and parameters separately: Now assume the original SDE has diagonal noise. Then, m = d and Jacobian matrix ∇σ i (z) can be written as: Consider the adjoint process for the augmented state along with the backward flow of the backward SDE. We write the overall state as X t = (Z t, (A z t), (A θ t) ), where we abuse notation slightly to let {Z t} t∈T denote the backward flow process. Then, by and, {X t} t∈T satisfies a backward SDE with a diffusion function that can be written as: Recall, for an SDE with diffusion function Σ(x) ∈ R d×m, it is said to satisfy the commutativity property (Rößler, 2004) When an SDE has commutative noise, the computationally intensive double Itô integrals (and the Lévy areas) need not be simulated by having the numerical scheme take advantage of the following property of iterated integrals : where the Brownian motion increment To see that the diffusion function indeed satisfies the commutativity condition, we consider several cases: • k = 1,..., d: Both LHS and RHS are zero unless j 1 = j 2 = k, since for Σ i,j 2 (x) to be non-zero, i = j 1 = j 2 = k. • k = d + 1..., 2d: Similar to the case above. • k = 2d + 1..., 2d + p:. Both LHS and RHS are zero unless Since in all scenarios, LHS = RHS, we conclude that the commutativity condition holds. Finally, we comment that the Milstein scheme for the stochastic adjoint of diagonal noise SDEs can be implemented such that during each iteration of the backward solve, vjp is only called a number of times constant with respect to the dimensionality of the original SDE. r e t u r n ans @staticmethod d e f backward (ctx, * grad_outputs): t s, flat_params_f, flat_params_g, * ans = c t x. s a v e d _ t e n s o r s f, g, dt, bm = c t x. f, c t x. g, c t x. dt, c t x. bm f_params, g_params = t u p l e (f . p a r a m e t e r s ), t u p l e (g . p a r a m e t e r s ) n _ t e n s o r s = l e n (ans) # Accumulate g r a d i e n t s a t i n t e r m e d i a t e p o i n t s. adj_y = _sequence_add (adj_y, t u p l e ( grad_outputs_ [ i − 1] f o r grad_outputs_ i n grad_outputs ) ) r e t u r n (* adj_y, None, None, None, adj_params_f, adj_params_g, None, None) In the following, α, β, and p are parameters of SDEs, and x 0 is a fixed initial value. Example 1. Analytical solution: Example 2. Analytical solution: Example 3. Analytical solution: In each numerical experiment, we duplicate the equation 10 times to obtain a system of SDEs where each dimension had their own parameter values sampled from the standard Gaussian distribution and then passed through a sigmoid to ensure positivity. Moreover, we also sample the initial value for each dimension from a Gaussian distribution. We consider training latent SDE models with our adjoint framework to recover a 1D Geometric Brownian motion, and a 3D stochastic Lorenz attractor process. The main objective is to verify that the learned posterior is able to reconstruct the training data, and the learned prior exhibit stochastic behavior. We jointly optimize the variational free energy with respect to parameters of the prior and posterior distributions at the initial latent state z 0, the prior and posterior drift, the diffusion function, the encoder, and the decoder. We include the details of dataset and architecture in Appendix N.1. For the stochastic Lorenz attractor, not only is the model able to reconstruct the data well, but also the learned prior process can produce bimodal samples in both data and latent space. This is showcased in the last row of Figure 4, where once the initial position sampled from the learned prior distribution is fixed, the latent and data space samples cluster around two modes. Note that this cannot be achieved by a latent ODE, where trajectories are determined once their initial latent state is determined. See Figure 4 for additional visualization on the synthetic Lorenz attractor dataset. See Figure 5 for visualization on the synthetic geometric Brownian motion dataset. We comment that for the second example, the posterior reconstructs the data well, and the prior process exhibit behavior of the data. However, from the third row, we can observe that the prior process is learned such that most of the uncertainty is account for in the initial latent state. We leave the investigation of more interpretable prior process for future work. Consider a geometric Brownian motion SDE: We use µ = 1, σ = 0.5, and x 0 = 0.1 + as the ground-truth model, where ∼ N(0, 0.03 2). We sample 1024 time series, each of which is observed at intervals of 0.02 from time 0 to time 1. We corrupt this data using Gaussian noise with mean zero and standard deviation 0.01. To recover the dynamics, we use a latent SDE model where the GRU has 1 layer and 100 hidden units, the prior and posterior drift functions are MLPs with 1 hidden layer of 100 units, and the diffusion function is an MLP with 1 hidden layer of 100 hidden units and the sigmoid activation applied at the end. The drift function in the posterior is time-inhomogenous in the sense that it takes in a context vector of size 1 at each observation that is output by the GRU from running backwards after processing all future observations. The decoder is a linear mapping from a 4 dimensional latent space to observation space. For all nonlinearities, we use the softplus function. We fix the observation model to be Gaussian with noise standard deviation 0.01. We optimize the model jointly with respect to the parameters of a Gaussian distribution for initial latent state distribution, the prior and posterior drift functions, the diffusion function, the GRU encoder, and the decoder. We use a fixed discretization with step size of 0.01 in both the forward and backward pass. We use the Adam optimizer with an initial learning rate of 0.01 that is decay by a factor of 0.999 after each iteration. We use a linear KL annealing schedule over the first 50 iterations. Consider a stochastic Lorenz attractor SDE with diagonal noise: dX t =σ (Y t − X t) dt + α x dW t, X 0 = x 0, dY t = (X t (ρ − Z t) − Y t ) dt + α y dW t, Y 0 = y 0, dZ t = (X t Y t − βZ t) dt + α z dW t, Z 0 = z 0. We use σ = 10, ρ = 28, β = 8/3, (α x, α y, α z) = (.1, .28., .3), and (x 0, y 0, z 0) sampled from the standard Gaussian distribution as the ground-truth model. We sample 1024 time series, each of which is observed at intervals of 0.025 from time 0 to time 1. We normalize these samples by their mean and standard deviation across each dimension and corrupt this data by Gaussian noise with mean zero and standard deviation 0.01. We use the same architecture and training procedure for the latent SDE model as in the geometric Brownian motion section, except that the diffusion function consists of four small neural networks, each for a single dimension of the latent SDE. We follow the preprocessing used by. Following Yıldız et al., we use a fully connected network to encode the first three observations of each sequence and thereafter predicted the remaining sequence. Note the current choice of encoder is for comparing fairly to models in the existing literature, and it may be extended to be a recurrent or attention model to enhance performance. The overall architecture is described in Appendix O and is similar to that of ODE 2 VAE Yıldız et al. with a similar number of parameters. We also use a fixed step size that is 1/5 of smallest interval between any two observations Yıldız et al.. We train latent ODE and latent SDE models with the Adam optimizer and its default hyperparameter settings, with an initial learning rate of 0.01 that is exponentially decayed with rate 0.999 during each iteration. We perform validation over the number of training iterations, KL penalty , and KL annealing schedule. All models were trained for at most 400 iterations, where we start to observe severe overfitting for most model instances. We use a latent SDE model with an MLP encoder which takes in the first three frames and outputs the mean and log-variance of the variational distribution of the initial latent state and a context vector. The decoder has a similar architecture as that for the ODE 2 VAE model Yıldız et al. and projects the 6-dimensional latent state into the 50-dimensional observation space. The posterior drift function takes in a 3-dimensional context vector output by the encoder and the current state and time, whereas the prior drift only takes in the current state and time. The diffusion function is composed of multiple small neural nets, each producing a scalar for the corresponding dimension such that the posterior SDE has diagonal noise. We comment that the overall parameter count of our model is smaller than that of ODE 2 VAE for the same task. The latent ODE baseline was implemented with a similar architecture, except is does not have the diffusion and prior drift components, and its vector field defining the ODE does not First row from left to right are the encoder and decoder. Second row from left to right are the prior drift, posterior drift, and diffusion functions. take in a context vector. Therefore, the model has slightly fewer parameters than the latent SDE model. See Figure 6 for overall details of the architecture. The main hyperparameter we tuned was the coefficient for reweighting the KL. For both the latent ODE and SDE, we considered training the model with a reweighting coefficient in {1, 0.1, 0.01, 0.001}, either with or without a linear KL annealing schedule that increased from 0 to the prescribed value over the first 200 iterations of training.

We present a constant memory gradient computation procedure through solutions of stochastic differential equations (SDEs) and apply the method for learning latent SDE models.

Represent the TLDR for finding the corresponding scientific article

We show shortcut connections should be placed in patterns that minimize between-layer distances during backpropagation, and design networks that achieve log L distances using L log(L) connections.