In this paper, we propose a new framework for training vision-based agent for First-Person Shooter (FPS) Game, in particular Doom. Our framework combines the state-of-the-art reinforcement learning approach (Asynchronous Advantage Actor-Critic (A3C) model [Mnih et al. (2016)]) with curriculum learning. Our model is simple in design and only uses game states from the AI side, rather than using opponents’ information [Lample & Chaplot (2016)]. On a known map, our agent won 10 out of the 11 attended games and the champion of Track1 in ViZDoom AI Competition 2016 by a large margin, 35% higher score than the second place.
1 INTRODUCTION
Deep Reinforcement Learning has achieved super-human performance in fully observable environments, e.g., in Atari Games [Mnih et al. (2015)] and Computer Go [Silver et al. (2016)]. Recently, Asynchronous Advantage Actor-Critic (A3C) [Mnih et al. (2016)] model shows good performance for 3D environment exploration, e.g. labyrinth exploration. However, in general, to train an agent in a partially observable 3D environment from raw frames remains an open challenge. Direct application of A3C to competitive 3D scenarios, e.g. 3D games, is nontrivial, partly due to sparse and long-term rewards in such scenarios.
Doom is a 1993 First-Person Shooter (FPS) game in which a player fights against other computercontrolled agents or human players in an adversarial 3D environment. Previous works on FPS AI [van Waveren (2001)] focused on using hand-tuned state machines and privileged information, e.g., the geometry of the map, the precise location of all players, to design playable agents. Although state-machine is conceptually simple and computationally efficient, it does not operate like human players, who only rely on visual (and possibly audio) inputs. Also, many complicated situations require manually-designed rules which could be time-consuming to tune.
In this paper, we train an AI agent in Doom with a framework that based on A3C with convolutional neural networks (CNN). This model uses only the recent 4 frames and game variables from the AI side, to predict the next action of the agent and the value of the current situation. We follow the curriculum learning paradigm [Bengio et al. (2009); Jiang et al. (2015)]: start from simple tasks and then gradually try harder ones. The difficulty of the task is controlled by a variety of parameters in Doom environment, including different types of maps, strength of the opponents and the design of the reward function. We also develop adaptive curriculum training that samples from a varying distribution of tasks to train the model, which is more stable and achieves higher score than A3C with the same number of epoch. As a result, our trained agent, named F1, won the champion in Track 1 of ViZDoom Competition 1 by a large margin.
There are many contemporary efforts on training a Doom AI based on the VizDoom platform [Kempka et al. (2016)] since its release. Arnold [Lample & Chaplot (2016)] also uses game frames and trains an action network using Deep Recurrent Q-learning [Hausknecht & Stone (2015)], and a navigation network with DQN [Mnih et al. (2015)]. However, there are several important differences. To predict the next action, they use a hybrid architecture (CNN+LSTM) that involves more complicated training procedure. Second, in addition to game frames, they require internal
1http://vizdoom.cs.put.edu.pl/competition-cig-2016/results
game status about the opponents as extra supervision during training, e.g., whether enemy is present in the current frame. IntelAct [Dosovitskiy & Koltun (2017)] models the Doom AI bot training in a supervised manner by predicting the future values of game variables (e.g., health, amount of ammo, etc) and acting accordingly. In comparison, we use curriculum learning with asynchronized actorcritic models and use stacked frames (4 most recent frames) and resized frames to mimic short-term memory and attention. Our approach requires no opponent’s information, and is thus suitable as a general framework to train agents for close-source games.
In VizDoom AI Competition 2016 at IEEE Computational Intelligence And Games (CIG) Conference2, our AI won the champion of Track1 (limited deathmatch with known map), and IntelAct won the champion of Track2 (full deathmatch with unknown maps). Neither of the two teams attends the other track. Arnold won the second places of both tracks and CLYDE [Ratcliffe et al. (2017)] won the third place of Track1.
2 THE ACTOR-CRITIC MODEL
The goal of Reinforcement Learning (RL) is to train an agent so that its behavior maximizes/minimizes expected future rewards/penalties it receives from a given environment [Sutton & Barto (1998)]. Two functions play important roles: a value function V (s) that gives the expected reward of the current state s, and a policy function π(a|s) that gives a probability distribution on the candidate actions a for the current state s. Getting the groundtruth value of either function would largely solve RL: the agent just follows π(a|s) to act, or jumps in the best state provided by V (s) when the number of candidate next states is finite and practically enumerable. However, neither is trivial.
Actor-critic models [Barto et al. (1983); Sutton (1984); Konda & Tsitsiklis (1999); Grondman et al. (2012)] aim to jointly estimate V (s) and π(a|s): from the current state st, the agent explores the environment by iteratively sampling the policy function π(at|st;wπ) and receives positive/negative reward, until the terminal state or a maximum number of iterations are reached. The exploration gives a trajectory {(st, at, rt), (st+1, at+1, rt+1), · · · }, from which the policy function and value function are updated. Specifically, to update the value function, we use the expected reward Rt along the trajectory as the ground truth; to update the policy function, we encourage actions that lead to high rewards, and penalize actions that lead to low rewards. To determine whether an action leads to high- or low-rewarding state, a reference point, called baseline [Williams (1992)], is usually needed. Using zero baseline might increase the estimation variance. [Peters & Schaal (2008)] gives a way to estimate the best baseline (a weighted sum of cumulative rewards) that minimizes the variance of the gradient estimation, in the scenario of episodic REINFORCE [Williams (1992)].
In actor-critic frameworks, we pick the baseline as the expected cumulative reward V (s) of the current state, which couples the two functions V (s) and π(a|s) together in the training, as shown in Fig. 1. Here the two functions reinforce each other: a correct π(a|s) gives high-rewarding trajectories which update V (s) towards the right direction; a correct V (s) picks out the correct actions for π(a|s) to reinforce. This mutual reinforcement behavior makes actor-critic model converge faster, but is also prone to converge to bad local minima, in particular for on-policy models that follow the very recent policy to sample trajectory during training. If the experience received by the agent in consecutive batches is highly correlated and biased towards a particular subset of the environment, then both π(a|s) and V (s) will be updated towards a biased direction and the agent may never see
2http://vizdoom.cs.put.edu.pl/competition-cig-2016
the whole picture. To reduce the correlation of game experience, Asynchronous Advantage ActorCritic Model [Mnih et al. (2016)] runs independent multiple threads of the game environment in parallel. These game instances are likely uncorrelated, therefore their experience in combination would be less biased.
For on-policy models, the same mutual reinforcement behavior will also lead to highly-peaked π(a|s) towards a few actions (or a few fixed action sequences), since it is always easy for both actor and critic to over-optimize on a small portion of the environment, and end up “living in their own realities”. To reduce the problem, [Mnih et al. (2016)] added an entropy term to the loss to encourage diversity, which we find to be critical. The final gradient update rules are listed as follows:
wπ ← wπ + α(Rt − V (st))∇wπ log π(at|st) + β∇wπH(π(·|st)) (1) wV ← wV − α∇wV (Rt − V (st)) 2 (2)
where Rt = ∑ T t′=t γt ′ −trt′ is the expected discounted reward at time t and α, β are the learning rate. In this work, we use Huber loss instead of the L2 loss in Eqn. 2.
Architecture. While [Mnih et al. (2016)] keeps a separate model for each asynchronous agent and perform model synchronization once in a while, we use an alternative approach called BatchA3C, in which all agents act on the same model and send batches to the main process for gradient descent optimization. The agents’ models are updated after each gradient update. Note that the contemporary work GA3C [Babaeizadeh et al. (2017)] also proposes a similar architecture. In their architecture, there is a prediction queue that collects agents’ experience and sends them to multiple predictors, and a training queue that collects experience to feed the optimization.
3 DOOM AS A REINFORCEMENT LEARNING PLATFORM
In Doom, the player controls the agent to fight against enemies in a 3D environment (e.g., in a maze). The agent can only see the environment from his viewpoint and thus receives partial information upon which it makes decisions. On modern computers, the original Doom runs in thousands of frames per second, making it suitable as a platform for training AI agent. ViZDoom [Kempka et al. (2016)] is an open-source platform that offers programming interface to communicate with Doom engine, ZDoom3. From the interface, users can obtain current frames of the game, and control the agent’s action. ViZDoom offers much flexibility, including:
Rich Scenarios. Many customized scenarios are made due to the popularity of the game, offering a variety of environments to train from. A scenario consists of many components, including 2D maps for the environment, scripts to control characters and events. Open-source tools, such as
3https://zdoom.org/
SLADE4, are also widely available to build new scenarios. We built our customized map (Fig. 2(b)) for training.
Game variables. In addition to image frames, ViZDoom environment also offers many games variables revealing the internal state of the game. This includes HEALTH, AMMO ? (agent’s health and ammunition), FRAG COUNT (current score) and so on. ViZDoom also offers USER? variables that are computed on the fly via scenario scripts. These USER? variables can provide more information of the agent, e.g., their spatial locations. Enemy information could also be obtained by modifying ViZDoom [Lample & Chaplot (2016)]. Such information is used to construct a reward function, or as a direct supervision to accelerate training [Lample & Chaplot (2016)].
Built-in bots. Built-in bots can be inserted in the battle. They are state machines with privileged information over the map and the player, which results in apparently decent intelligence with minimal computational cost. By competing against built-in bots, the agent learns to improve.
Evaluation Criterion. In FPS games, to evaluate their strength, multiple AIs are placed to a scenario for a deathmatch, in which every AI plays for itself against the remaining AIs. Frags per episode, the number of kills minus the number of suicides for the agent in one round of game, is often used as a metric. An AI is stronger if its frags is ranked higher against others. In this work, we use an episode of 2-minute game time (4200 frames in total) for all our evaluations unless noted otherwise.
4 METHOD
4.1 NETWORK ARCHITECTURE
We use convolutional neural networks to extract features from the game frames and then combine its output representation with game variables. Fig. 3 shows the network architecture and Tbl. 1 gives the parameters. It takes the frames as the input (i.e., the state s) and outputs two branches, one that outputs the value function V (s) by regression, while the other outputs the policy function π(s|a) by a regular softmax. The parameters of the two functions are shared before the branch.
For input, we use the most recent 4 frames plus the center part of them, scaled to the same size (120 × 120). Therefore, these centered “attention frames” have higher resolution than regular game frames, and greatly increase the aiming accuracy. The policy network will give 6 actions, namely MOVE FORWARD, MOVE LEFT, MOVE RIGHT, TURN LEFT, TURN RIGHT, and ATTACK. We found other on-off actions (e.g., MOVE BACKWARD) offered by ViZDoom less important. After feature extraction by convolutional network, game variables are incorporated. This includes the agent’s Health (0-100) and Ammo (how many bullets left). They are related to AI itself and thus legal in the game environment for training, testing and ViZDoom AI competition.
4.2 TRAINING PIPELINE
Our training procedure is implemented with TensorFlow [Abadi et al. (2016)] and tensorpack5. We open 255 processes, each running one Doom instance, and sending experience (st, at, rt) to the
4http://slade.mancubus.net/ 5https://github.com/ppwwyyxx/tensorpack
main process which runs the training procedure. The main process collects frames from different game instances to create batches, and optimizes on these batches asynchronously on one or more GPUs using Eqn. 1 and Eqn. 2. The frames from different processes running independent game instances, are likely to be uncorrelated, which stabilizes the training. This procedure is slightly different from the original A3C, where each game instance collects their own experience and updates the parameters asynchronously.
Despite the use of entropy term, we still find that π(·|s) is highly peaked. Therefore, during trajectory exploration, we encourage exploration by the following changes: a) multiply the policy output of the network by an exploration factor (0.2) before softmax b) uniformly randomize the action for 10% random frames.
As mentioned in [Kempka et al. (2016)], care should be taken for frame skips. Small frame skip introduces strong correlation in the training set, while big frame skip reduces effective training samples. We set frame skip to be 3. We choose 640x480 as the input frame resolution and do not use high aspect ratio resolution [Lample & Chaplot (2016)] to increase the field of view.
We use Adam [Kingma & Ba (2014)] with ǫ = 10−3 for training. Batch size is 128, discount factor γ = 0.99, learning rate α = 10−4 and the policy learning rate β = 0.08α. The model is trained from scratch. The training procedure runs on Intel Xeon CPU E5-2680v2 at 2. 80GHz, and 2 TitanX GPUs. It takes several days to obtain a decent result. Our final model, namely the F1 bot, is trained for around 3 million mini-batches on multiple different scenarios.
4.3 CURRICULUM LEARNING
When the environment only gives very sparse rewards, or adversarial, A3C takes a long time to converge to a satisfying solution. A direct training with A3C on the map CIGTrack1 with 8 builtin bots does not yield sensible performance. To address this, we use curriculum learning [Bengio et al. (2009)] that trains an agent with a sequence of progressively more difficult environments. By varying parameters in Doom (Sec. 3), we could control its difficulty level.
Reward Shaping. Reward shaping has been shown to be an effective technique to apply reinforcement learning in a complicated environment with delayed reward [Ng et al. (1999); Devlin et al. (2011)]. In our case, besides the basic reward for kills (+1) and death (-1), intermediate rewards are used as shown in Tbl. 2. We penalize agent with a living state, encouraging it to explore and encounter more enemies. health loss and ammo loss place linear reward for a decrement of health and ammunition. ammo pickup and health pickup place reward for picking up these two items. In addition, there is extra reward for picking up ammunition when in need (e.g. almost out of ammo). dist penalty and dist reward push the agent away from the previous locations, encouraging it to explore. The penalty is applied every action, when the displacement of the bot relative to the last state is less than a threshold dist penalty thres. And dist reward is applied for every unit displacement the agent makes. Similar to [Lample & Chaplot (2016)], the displacement information is computed from the ground truth location variables provided by Doom engine, and will not be used in the competition. However, unlike [Lample & Chaplot (2016)] that uses enemy-in-sight signal for training, locations can be extracted directly from USER? variables, or can easily be computed roughly with action history.
Curriculum Design. We train the bot on FlatMap that contains a simple square with a few pillars (Fig. 2(a)) with several curricula (Tbl. 3), and then proceed to CIGTrack1. For each map, we design curricula by varying the strength of built-in bots, i.e., their moving speed, initial health and initial weapon. Our agent always uses RocketLauncher as its only weapon. Training on FlatMap leads to a capable initial model which is quickly adapted to more complicated maps. As shown in Tbl. 2, for CIGTrack1 we increase dist penalty thres to keep the agent moving, and increase num bots so that the agent encounters more enemies per episode.
Adaptive Curriculum. In addition to staged curriculum learning, we also design adaptive curriculum learning by assigning a probability distribution on different levels for each thread that runs a Doom instance. The probability distribution shifts towards more difficult curriculum when the agent performs well on the current distribution, and shifts towards easier level otherwise. We consider the agent to perform well if its frag count is greater than 10 points.
4.4 POST-TRAINING RULES
For a better performance in the competition, we also put several rules to process the action given by the trained policy network, called post-training (PT) rules. There are two sets of buttons in ViZDoom: on-off buttons and delta buttons. While on-off button maps to the binary states of a keystroke (e.g., pressing the up arrow key will move the agent forward), delta buttons mimic the mouse behavior and could act faster in certain situations. Therefore, we setup rules that detect the intention of the agent and accelerate with delta button. For example, when the agent turns by invoking TURN LEFT repeatedly, we convert its action to TURN LEFT RIGHT DELTA for acceleration. Besides, the trained model might get stuck in rare situations, e.g., keep moving forward but blocked by an explosive bucket. We also designed rules to detect and fix them.
5 EXPERIMENT
In this section, we show the training procedure (Sec. 5.1), evaluate our AIs with ablation analysis (Sec. 5.2) and ViZDoom AI Competition (Sec. 5.3). We mainly compare among three AIs: (1) F1Pre, the bot trained with FlatMap only, (2) F1Plain, the bot trained on both FlatMap and CIGTrack1, but without post-training rules, and (3) the final F1 bot that attends competition.
5.1 CURRICULUM LEARNING ON FLATMAP
Fig. 4 shows that the curriculum learning increases the performance of the agents over all levels. When an agent becomes stronger in the higher level of class, it is also stronger in the lower level of class without overfitting. Fig. 5 shows comparison between adaptive curriculum learning with pure A3C. We can see that pure A3C can learn on FlatMap but is slower. Moreover, in CIGTrack1, a direct application of A3C does not yield sensible performance.
5.2 ABLATION ANALYSIS
Visualization. Fig. 6 shows the visualization of the first convolutional layer of the trained AI agent. We could see that the convolutional kernels of the current frame is less noisy than the kernels of previous frames. This means that the agent makes the most use of the current frames.
Effect of History Frames. Interestingly, while the agent focuses on the current frame, it also uses motion information. For this, we use (1) 4 duplicated current frames (2) 4 recent frames in reverse order, as the input. This gives 8.50 and 2.39 mean frags, compared to 10.34 in the normal case, showing that the agent heavily uses the motion information for better decision. In particular, the bot is totally confused with the reversed motion feature. Detailed results are shown in Tbl. 5.
Post-training Rules. Tbl. 5 shows that the post-training rules improve the performance. As a future work, an end-to-end training involving delta buttons could make the bot better.
Internal Tournament. We also evaluate our AIs with internal tournaments (Tbl. 4). All our bots beat the performance of built-in bots by a large margin, even though they use privileged information. F1Pre, trained with only FlatMap, shows decent performance, but is not as good as the models trained with both FlatMap and CIGTrack1. The final bot F1 performs the best.
Behaviors. Visually, the three bots behave differently. F1Pre is a bit overtrained in FlatMap and does not move too often, but when it sees enemies, even faraway, it will start to shoot. Occasionally it will move to the corner and pick medkits. In CIGTrack1, F1Pre stays in one place and ambushes opponents who pass by. On the other hand, F1Plain and F1 always move forwards and turn at the corner. As expected, F1 moves and turns faster.
Tactics All bots develop interesting local tactics when exchanging fire with enemy: they slide around when shooting the enemy. This is quite effective for dodging others’ attack. Also when they shoot the enemy, they usually take advantage of the splashing effect of rocket to cause additional damage for enemy, e.g., shooting the wall when the enemy is moving. They do not pick ammunition too often, even if they can no longer shoot. However, such disadvantage is mitigated by the nature of deathmatch: when a player dies, it will respawn with ammunition. We also check states with highest/lowest estimated future value V (s) over a 10-episode evaluation of F1 bot, from which we can speculate its tactics. The highest value is V = 0.97 when the agent fired, and about to hit the enemy. One low value is V = −0.44, ammo = 0, when the agent encountered an enemy at the corner but is out of ammunition. Both cases are reasonable.
5.3 COMPETITION
We attended the ViZDoom AI Competition hosted by IEEE CIG. There are 2 tracks in the competition. Track 1 (Limited Deathmatch) uses a known map and fixed weapons, while Track 2 (Full Deathmatch) uses 3 unknown maps and a variety of weapons. Each bot fights against all others for 12 rounds of 10 minutes each. Due to server capacity, each bot skips one match in the first 9 rounds. All bots are supposed to run in real-time (>35fps) on a GTX960 GPU.
Our F1 bot won 10 out of 11 attended games and won the champion for Track 1 by a large margin. We have achieved 559 frags, 35.4% higher than 413 frags achieved by Arnold [Lample & Chaplot (2016)], that uses extra game state for model training. On the other hand, IntelAct [Dosovitskiy & Koltun (2017)] won Track 2. The full videos for the two tracks have been released67, as well as an additional game between Human and AIs8. Our bot behaves reasonable and very human-like in Track 1. In the match between Human and AIs, our bot was even ahead of the human player for a short period (6:30 to 7:00).
6 CONCLUSION
Teaching agents to act properly in complicated and adversarial 3D environment is a very challenging task. In this paper, we propose a new framework to train a strong AI agent in a First-Person Shooter (FPS) game, Doom, using a combination of state-of-the-art Deep Reinforcement Learning and Curriculum Training. Via playing against built-in bots in a progressive manner, our bot wins the champion of Track1 (known map) in ViZDoom AI Competition. Furthermore, it learns to use motion features and build its own tactics during the game, which is never taught explicitly.
Currently, our bot is still an reactive agent that only remembers the last 4 frames to act. Ideally, a bot should be able to build a map from an unknown environment and localize itself, is able to have a global plan to act, and visualize its reasoning process. We leave them to future works.

The paper describes approaches taken to train learning agents for the 3D game Doom. The authors propose a number of performance enhancements (curriculum learning, attention (zoomed-in centered) frames, reward shaping, game variables, post-training rules) inspired by domain knowledge.
The enhancements together lead to a clear win as demonstrated by the competition results. From Fig 4, the curriculum learning clearly helps with learning over increasingly difficult settings. A nice result is that there is no overfitting to the harder classes once they have learned (probably because the curriculum is health and speed). The authors conclude from Fig 5 that the adaptive curriculum is better and more stable that pure A3C; however, this is a bit of a stretch given that graph. They go on to say that Pure A3C doesn't learn at all in the harder map but then show no result/graph to back this claim. Tbl 5 shows a clear benefit of the post-training rules.
If the goal is to solve problems like these (3D shooters), then this paper makes a significant contribution in that it shows which techniques are practical for solving the problem and ultimately improving performance in these kinds of tasks. Still, I am just not excited about this paper, mainly because it relies so heavily of many sources of domain knowledge, it is quite far from the pure reinforcement learning problem. The results are relatively unsurprising. Maybe they are novel for this problem, though.
I'm not sure we can realistically draw any conclusions about Figure 6 in the paper's current form. I recommend the authors increase the resolution or run some actual metrics to determine the fuzziness/clarity of each row/image: something more concrete than an arrow of already low-resolution images.
--- Added after rebuttal:
I still do not see any high-res images for Figure 6 or any link to them, but I trust that the authors will add them if accepted.


Ensembles of neural networks are known to be much more robust and accurate than individual networks. However, training multiple deep networks for model averaging is computationally expensive. In this paper, we propose a method to obtain the seemingly contradictory goal of ensembling multiple neural networks at no additional training cost. We achieve this goal by training a single neural network, converging to several local minima along its optimization path and saving the model parameters. To obtain repeated rapid convergence, we leverage recent work on cyclic learning rate schedules. The resulting technique, which we refer to as Snapshot Ensembling, is simple, yet surprisingly effective. We show in a series of experiments that our approach is compatible with diverse network architectures and learning tasks. It consistently yields lower error rates than state-of-the-art single models at no additional training cost, and compares favorably with traditional network ensembles. On CIFAR-10 and CIFAR-100 our DenseNet Snapshot Ensembles obtain error rates of 3.4% and 17.4% respectively.
1 INTRODUCTION
Stochastic Gradient Descent (SGD) (Bottou, 2010) and its accelerated variants (Kingma & Ba, 2014; Duchi et al., 2011) have become the de-facto approaches for optimizing deep neural networks. The popularity of SGD can be attributed to its ability to avoid and even escape spurious saddle-points and local minima (Dauphin et al., 2014). Although avoiding these spurious solutions is generally considered positive, in this paper we argue that these local minima contain useful information that may in fact improve model performance.
Although deep networks typically never converge to a global minimum, there is a notion of “good” and “bad” local minima with respect to generalization. Keskar et al. (2016) argue that local minima with flat basins tend to generalize better. SGD tends to avoid sharper local minima because gradients are computed from small mini-batches and are therefore inexact (Keskar et al., 2016). If the learningrate is sufficiently large, the intrinsic random motion across gradient steps prevents the optimizer from reaching any of the sharp basins along its optimization path. However, if the learning rate is small, the model tends to converge into the closest local minimum. These two very different behaviors of SGD are typically exploited in different phases of optimization (He et al., 2016a). Initially the learning rate is kept high to move into the general vicinity of a flat local minimum. Once this search has reached a stage in which no further progress is made, the learning rate is dropped (once or twice), triggering a descent, and ultimately convergence, to the final local minimum.
It is well established (Kawaguchi, 2016) that the number of possible local minima grows exponentially with the number of parameters—of which modern neural networks can have millions. It is therefore not surprising that two identical architectures optimized with different initializations or minibatch orderings will converge to different solutions. Although different local minima often have very similar error rates, the corresponding neural networks tend to make different mistakes. This
∗Authors contribute equally.
diversity can be exploited through ensembling, in which multiple neural networks are trained from different initializations and then combined with majority voting or averaging (Caruana et al., 2004). Ensembling often leads to drastic reductions in error rates. In fact, most high profile competitions, e.g. Imagenet (Deng et al., 2009) or Kaggle1, are won by ensembles of deep learning architectures.
Despite its obvious advantages, the use of ensembling for deep networks is not nearly as widespread as it is for other algorithms. One likely reason for this lack of adaptation may be the cost of learning multiple neural networks. Training deep networks can last for weeks, even on high performance hardware with GPU acceleration. As the training cost for ensembles increases linearly, ensembles can quickly becomes uneconomical for most researchers without access to industrial scale computational resources.
In this paper we focus on the seemingly-contradictory goal of learning an ensemble of multiple neural networks without incurring any additional training costs. We achieve this goal with a training method that is simple and straight-forward to implement. Our approach leverages the non-convex nature of neural networks and the ability of SGD to converge to and escape from local minima on demand. Instead of training M neural networks independently from scratch, we let SGD converge M times to local minima along its optimization path. Each time the model converges, we save the weights and add the corresponding network to our ensemble. We then restart the optimization with a large learning rate to escape the current local minimum. More specifically, we adopt the cycling procedure suggested by Loshchilov & Hutter (2016), in which the learning rate is abruptly raised and then quickly lowered to follow a cosine function. Because our final ensemble consists of snapshots of the optimization path, we refer to our approach as Snapshot Ensembling. Figure 1 presents a high-level overview of this method.
In contrast to traditional ensembles, the training time for the entire ensemble is identical to the time required to train a single traditional model. During testing time, one can evaluate and average the last (and therefore most accurate) m out of M models. Our approach is naturally compatible with other methods to improve the accuracy, such as data augmentation, stochastic depth (Huang et al., 2016b), or batch normalization (Ioffe & Szegedy, 2015). In fact, Snapshot Ensembles can even be ensembled, if for example parallel resources are available during training. In this case, an ensemble of K Snapshot Ensembles yields K ×M models at K times the training cost. We evaluate the efficacy of Snapshot Ensembles on three state-of-the-art deep learning architectures for object recognition: ResNet (He et al., 2016b), Wide-ResNet (Zagoruyko & Komodakis, 2016), and DenseNet (Huang et al., 2016a). We show across four different data sets that Snapshot Ensembles almost always reduce error without increasing training costs. For example, on CIFAR-10 and CIFAR-100, Snapshot Ensembles obtains error rates of 3.44% and 17.41% respectively.
1www.kaggle.com
2 RELATED WORK
Neural network ensembles have been widely studied and applied in machine learning (Hansen & Salamon, 1990; Krogh et al., 1995). However, most of these prior studies focus on improving the generalization performance, while few of them address the cost of training ensembles.
As an alternative to traditional ensembles, so-called “implicit” ensembles have high efficiency during both training and testing (Srivastava et al., 2014; Wan et al., 2013; Huang et al., 2016b; Singh et al., 2016; Krueger et al., 2016). The Dropout (Srivastava et al., 2014) technique creates an ensemble out of a single model by “dropping” — or zeroing — random sets of hidden nodes during each mini-batch. At test time, no nodes are dropped, and each node is scaled by the probability of surviving during training. Srivastava et al. claim that Dropout reduces overfitting by preventing the co-adaptation of nodes. An alternative explanation is that this mechanism creates an exponential number of networks with shared weights during training, which are then implicitly ensembled at test time. DropConnect (Wan et al., 2013) uses a similar trick to create ensembles at test time by dropping connections (weights) during training instead of nodes. The recently proposed Stochastic Depth technique (Huang et al., 2016b) randomly drops layers during training to create an implicit ensemble of networks with varying depth at test time. Finally, Swapout (Singh et al., 2016) is a stochastic training method that generalizes Dropout and Stochastic Depth. From the perspective of model ensembling, Swapout creates diversified network structures for model averaging. Our proposed method similarly trains only a single model; however, the resulting ensemble is “explicit” in that the models do not share weights. Furthermore, our method can be used in conjunction with any of these implicit ensembling techniques.
Several recent publications focus on reducing the test time cost of ensembles, by transferring the “knowledge” of cumbersome ensembles into a single model (Bucilu et al., 2006; Hinton et al., 2015). Hinton et al. (2015) propose to use an ensemble of multiple networks as the target of a single (smaller) network. Our proposed method is complementary to these works as we aim to reduce the training cost of ensembles rather than the test-time cost.
Perhaps most similar to our work is that of Swann & Allinson (1998) and Xie et al. (2013), who explore creating ensembles from slices of the learning trajectory. Xie et al. introduce the horizontal and vertical ensembling method, which combines the output of networks within a range of training epochs. More recently, Jean et al. (2014) and Sennrich et al. (2016) show improvement by ensembling the intermediate stages of model training. Laine & Aila (2016) propose a temporal ensembling method for semi-supervised learning, which achieves consensus among models trained with different regularization and augmentation conditions for better generalization performance. Finally, Moghimi et al. (2016) show that boosting can be applied to convolutional neural networks to create strong ensembles. Our work differs from these prior works in that we force the model to visit multiple local minima, and we take snapshots only when the model reaches a minimum. We believe this key insight allows us to leverage more power from our ensembles.
Our work is inspired by the recent findings of Loshchilov & Hutter (2016) and Smith (2016), who show that cyclic learning rates can be effective for training convolutional neural networks. The authors show that each cycle produces models which are (almost) competitive to those learned with traditional learning rate schedules while requiring a fraction of training iterations. Although model performance temporarily suffers when the learning rate cycle is restarted, the performance eventually surpasses the previous cycle after annealing the learning rate. The authors suggest that cycling perturbs the parameters of a converged model, which allows the model to find a better local minimum. We build upon these recent findings by (1) showing that there is significant diversity in the local minima visited during each cycle and (2) exploiting this diversity using ensembles. We are not concerned with speeding up or improving the training of a single model; rather, our goal is to extract an ensemble of classifiers while following the optimization path of the final model.
3 SNAPSHOT ENSEMBLING
Snapshot Ensembling produces an ensemble of accurate and diverse models from a single training process. At the heart of Snapshot Ensembling is an optimization process which visits several local minima before converging to a final solution. We take model snapshots at these various minima, and average their predictions at test time.
Ensembles work best if the individual models (1) have low test error and (2) do not overlap in the set of examples they misclassify. Along most of the optimization path, the weight assignments of a neural network tend not to correspond to low test error. In fact, it is commonly observed that the validation error drops significantly only after the learning rate has been reduced, which is typically done after several hundred epochs. Our approach is inspired by the observation that training neural networks for fewer epochs and dropping the learning rate earlier has minor impact on the final test error (Loshchilov & Hutter, 2016). This seems to suggest that local minima along the optimization path become promising (in terms of generalization error) after only a few epochs.
Cyclic Cosine Annealing. To converge to multiple local minima, we follow a cyclic annealing schedule as proposed by Loshchilov & Hutter (2016). We lower the learning rate at a very fast pace, encouraging the model to converge towards its first local minimum after as few as 50 epochs. The optimization is then continued at a larger learning rate, which perturbs the model and dislodges it from the minimum. We repeat this process several times to obtain multiple convergences. Formally, the learning rate α has the form:
α(t) = f (mod (t− 1, dT/Me)) , (1)
where t is the iteration number, T is the total number of training iterations, and f is a monotonically decreasing function. In other words, we split the training process into M cycles, each of which starts with a large learning rate, which is annealed to a smaller learning rate. The large learning rate α = f(0) provides the model enough energy to escape from a critical point, while the small learning rate
α = f(dT/Me) drives the model to a well behaved local minimum. In our experiments, we set f to be the shifted cosine function proposed by Loshchilov & Hutter (2016):
α(t) = α0 2
( cos ( πmod(t− 1, dT/Me)
dT/Me
) + 1 ) , (2)
where α0 is the initial learning rate. Intuitively, this function anneals the learning rate from its initial value α0 to f(dT/Me) ≈ 0 over the course of a cycle. Following (Loshchilov & Hutter, 2016), we update the learning rate at each iteration rather than at every epoch. This improves the convergence of short cycles, even when a large initial learning rate is used. Snapshot Ensembling. Figure 2 depicts the training process using cyclic and traditional learning rate schedules. At the end of each training cycle, it is apparent that the model reaches a local minimum with respect to the training loss. Thus, before raising the learning rate, we take a “snapshot” of the model weights (indicated as vertical dashed black lines). After training M cycles, we have M model snapshots, f1 . . . fM , each of which will be used in the final ensemble. It is important to highlight that the total training time of the M snapshots is the same as training a model with a standard schedule (indicated in blue). In some cases, the standard learning rate schedule achieves lower training loss than the cyclic schedule; however, as we will show in the next section, the benefits of ensembling outweigh this difference. Ensembling at Test Time. The ensemble prediction at test time is the average of the last m (m ≤ M ) model’s softmax outputs. Let x be a test sample and let hi (x) be the softmax score of snapshot i. The output of the ensemble is a simple average of the last m models: hEnsemble = 1 m ∑m−1 0 hM−i (x) . We always ensemble the last m models, as these models tend to have the lowest test error.
4 EXPERIMENTS
We demonstrate the effectiveness of Snapshot Ensembles on several benchmark datasets, comparing with competitive baselines. We run all experiments with Torch 7 (Collobert et al., 2011)2.
4.1 DATASETS
CIFAR. The two CIFAR datasets (Krizhevsky & Hinton, 2009) consist of colored natural images sized at 32×32 pixels. CIFAR-10 (C10) and CIFAR-100 (C100) images are drawn from 10 and 100 classes, respectively. For each dataset, there are 50,000 training images and 10,000 images reserved for testing. We use a standard data augmentation scheme (Lin et al., 2013; Romero et al., 2014; Lee et al., 2015; Springenberg et al., 2014; Srivastava et al., 2015; Huang et al., 2016b; Larsson et al., 2016), in which the images are zero-padded with 4 pixels on each side, randomly cropped to produce 32×32 images, and horizontally mirrored with probability 0.5. SVHN. The Street View House Numbers (SVHN) dataset (Netzer et al., 2011) contains 32 × 32 colored digit images from Google Street View, with one class for each digit. There are 73,257 images in the training set and 26,032 images in the test set. Following common practice (Sermanet et al., 2012; Goodfellow et al., 2013; Huang et al., 2016a), we withhold 6,000 training images for validation, and train on the remaining images without data augmentation. Tiny ImageNet. The Tiny ImageNet dataset3 consists of a subset of ImageNet images (Deng et al., 2009). There are 200 classes, each of which has 500 training images and 50 validation images. Each image is resized to 64×64 and augmented with random crops, horizontal mirroring, and RGB intensity scaling (Krizhevsky et al., 2012). ImageNet. The ILSVRC 2012 classification dataset (Deng et al., 2009) consists of 1000 images classes, with a total of 1.2 million training images and 50,000 validation images. We adopt the same
2Code to reproduce results is available at https://github.com/gaohuang/SnapshotEnsemble 3https://tiny-imagenet.herokuapp.com
data augmentation scheme as in (He et al., 2016a; Huang et al., 2016a) and apply a 224×224 center crop to images at test time.
4.2 TRAINING SETTING
Architectures. We test several state-of-the-art architectures, including residual networks (ResNet) (He et al., 2016a), Wide ResNet (Zagoruyko & Komodakis, 2016) and DenseNet (Huang et al., 2016a). For ResNet, we use the original 110-layer network introduced by He et al. (2016a). Wide-ResNet is a 32-layer ResNet with 4 times as many convolutional features per layer as a standard ResNet. For DenseNet, our large model follows the same setup as (Huang et al., 2016a), with depth L = 100 and growth rate k = 24. In addition, we also evaluate our method on a small DenseNet, with depth L = 40 and k = 12. To adapt all these networks to Tiny ImageNet, we add a stride of 2 to the first layer of the models, which downsamples the images to 32×32. For ImageNet, we test the 50-layer ResNet proposed in (He et al., 2016a). We use a mini batch size of 64.4
Baselines. Snapshot Ensembles incur the training cost of a single model; therefore, we compare with baselines that require the same amount of training. First, we compare against a Single Model trained with a standard learning rate schedule, dropping the learning rate from 0.1 to 0.01 halfway through training, and then to 0.001 when training is at 75%. Additionally, to compare against implicit ensembling methods, we test against a single model trained with Dropout. This baseline uses the same learning rate as above, and drops nodes during training with a probability of 0.2.
We then test the Snapshot Ensemble algorithm trained with the cyclic cosine learning rate as described in (2). We test models with the max learning rate α0 set to 0.1 and 0.2. In both cases, we divide the training process into learning rate cycles. Model snapshots are taken after each learning rate cycle. Additionally, we train a Snapshot Ensemble with a non-cyclic learning rate schedule. This NoCycle Snapshot Ensemble, which uses the same schedule as the Single Model and Dropout baselines, is meant to highlight the impact of cyclic learning rates for our method. To accurately compare with the cyclic Snapshot Ensembles, we take the same number of snapshots equally spaced throughout the training process. Finally, we compare against SingleCycle Ensembles, a Snapshot Ensemble variant in which the network is re-initialized at the beginning of every cosine learning rate cycle, rather than using the parameters from the previous optimization cycle. This baseline essentially creates a traditional ensemble, yet each network only has 1/M of the typical training time. This variant is meant to highlight the tradeoff between model diversity and model convergence. Though SingleCycle Ensembles should in theory explore more of the parameter space, the models do not benefit from the optimization of previous cycles. Training Budget. On CIFAR datasets, the training budget is B = 300 epochs for DenseNet-40 and DenseNet-100, and B = 200 for ResNet and Wide ResNet models. Snapshot variants are trained with M = 6 cycles of B/M = 50 epochs for DenseNets, and M = 5 cycles of B/M = 40 epochs for ResNets/Wide ResNets. SVHN models are trained with a budget of B = 40 epochs (5 cycles of 8 epochs). For Tiny ImageNet, we use a training budget of B = 150 (6 cycles of 25 epochs). Finally, ImageNet is trained with a budget of B = 90 epochs, and we trained 2 Snapshot variants: one with M = 2 cycles and one with M = 3.
4Exceptions: ResNet-110 and Wide-ResNet are trained with batch size 128 on Tiny ImageNet. The ImageNet model is trained with batch size 256.
4.3 SNAPSHOT ENSEMBLE RESULTS
Method Val. Error (%)
Single model 24.01 Snapshot Ensemble (M = 2) 23.33 Snapshot Ensemble (M = 3) 23.96
Table 2: Top-1 error rates (%) on ImageNet validation set using ResNet-50 with varying number of cycles.
Accuracy. The main results are summarized in Table 1. In most cases, Snapshot ensembles achieve lower error than any of the baseline methods. Most notably, Snapshot Ensembles yield an error rate of 17.41% on CIFAR100 using large DenseNets, far outperforming the record of 19.25% under the same training cost and architecture (Huang et al., 2016a). Our method has the most success on CIFAR-100 and Tiny ImageNet, which is likely due to the
complexity of these datasets. The softmax outputs for these datasets are high dimensional due to the large number of classes, making it unlikely that any two models make the same predictions. Snapshot Ensembling is also capable of improving the competitive baselines for CIFAR-10 and SVHN as well, reducing error by 1% and 0.4% respectively with the Wide ResNet architecture.
The NoCycle Snapshot Ensemble generally has little effect on performance, and in some instances even increases the test error. This highlights the need for a cyclic learning rate for useful ensembling. The SingleCycle Ensemble has similarly mixed performance. In some cases, e.g., DenseNet-40 on CIFAR-100, the SingleCycle Ensemble is competitive with Snapshot Ensembles. However, as the model size increases to 100 layers, it does not perform as well. This is because it is difficult to train a large model from scratch in only a few epochs. These results demonstrate that Snapshot Ensembles tend to work best when utilizing information from previous cycles. Effectively, Snapshot Ensembles strike a balance between model diversity and optimization.
Table 2 shows Snapshot Ensemble results on ImageNet. The Snapshot Ensemble with M = 2 achieves 23.33% validation error, outperforming the single model baseline with 24.01% validation error. It appears that 2 cycles is the optimal choice for the ImageNet dataset. Provided with the limited total training budget B = 90 epochs, we hypothesize that allocating fewer than B/2 = 45 epochs per training cycle is insufficient for the model to converge on such a large dataset.
Ensemble Size. In some applications, it may be beneficial to vary the size of the ensemble dynamically at test time depending on available resources. Figure 3 displays the performance of DenseNet-40 on the CIFAR-100 dataset as the effective ensemble size,m, is varied. Each ensemble consists of snapshots from later cycles, as these snapshots have received the most training and therefore have likely converged to better minima. Although ensembling more models generally gives better performance, we observe significant drops in error when the second and third models are added to the ensemble. In most cases, an ensemble of two models outperforms the baseline model. Restart Learning Rate. The effect of the restart learning rate can be observed in Figure 3. The left two plots show performance when using a restart learning rate of α0 = 0.1 at the beginning of each cycle, and the right two plots show α0 = 0.2. In most cases, ensembles with the larger restart learning rate perform better, presumably because the strong perturbation in between cycles increases the diversity of local minima.
Varying Number of Cycles. Given a fixed training budget, there is a trade-off between the number of learning rate cycles and their length. Therefore, we investigate how the number of cycles M affects the ensemble performance, given a fixed training budget. We train a 40-layer DenseNet on the CIFAR-100 dataset with an initial learning rate of α0 = 0.2. We fix the total training budget B = 300 epochs, and vary the value of M ∈ {2, 4, 6, 8, 10}. As shown in Table 3, our method is relatively robust with respect to different values of M . At the extremes, M = 2 and M = 10, we find a slight degradation in performance, as the cycles are either too few or too short. In practice, we find that setting M to be 4 ∼ 8 works reasonably well. Varying Training Budget. The left and middle panels of Figure 4 show the performance of Snapshot Ensembles and SingleCycle Ensembles as a function of training budget (where the number of cycles is fixed at M = 6). We train a 40-layer DenseNet on CIFAR-10 and CIFAR-100, with an initial learning rate of α0 = 0.1, varying the total number of training epochs from 60 to 300. We observe
that both Snapshot Ensembles and SingleCycle Ensembles become more accurate as training budget increases. However, we note that as training budget decreases, Snapshot Ensembles still yield competitive results, while the performance of the SingleCycle Ensembles degrades rapidly. These results highlight the improvements that Snapshot Ensembles obtain when the budget is low. If the budget is high, then the SingleCycle baseline approaches true ensembles and outperforms Snapshot ensembles eventually. Comparison with True Ensembles. We compare Snapshot Ensembles with the traditional ensembling method. The right panel of Figure 4 shows the test error rates of DenseNet-40 on CIFAR-100. The true ensemble method averages models that are trained with 300 full epochs, each with different weight initializations. Given the same number of models at test time, the error rate of the true ensemble can be seen as a lower bound of our method. Our method achieves performance that is comparable with ensembling of 2 independent models, but with the training cost of one model.
4.4 DIVERSITY OF MODEL ENSEMBLES
Parameter Space. We hypothesize that the cyclic learning rate schedule creates snapshots which are not only accurate but also diverse with respect to model predictions. We qualitatively measure this diversity by visualizing the local minima that models converge to. To do so, we linearly interpolate snapshot models, as described by Goodfellow et al. (2014). Let J (θ) be the test error of a model using parameters θ. Given θ1 and θ2 — the parameters from models 1 and 2 respectively — we can compute the loss for a convex combination of model parameters: J (λ (θ1) + (1− λ) (θ2)), where λ is a mixing coefficient. Setting λ to 1 results in a parameters that are entirely θ1 while setting λ to 0 gives the parameters θ2. By sweeping the values of λ, we can examine a linear slice of the parameter space. Two models that converge to a similar minimum will have smooth parameter interpolations, whereas models that converge to different minima will likely have a non-convex interpolation, with a spike in error when λ is between 0 and 1.
Figure 5 displays interpolations between the final model of DenseNet-40 (sixth snapshot) and all intermediate snapshots. The left two plots show Snapshot Ensemble models trained with a cyclic learning rate, while the right two plots show NoCycle Snapshot models. λ = 0 represents a model which is entirely snapshot parameters, while λ = 1 represents a model which is entirely the parameters of the final model. From this figure, it is clear that there are differences between cyclic and
non-cyclic learning rate schedules. Firstly, all of the cyclic snapshots achieve roughly the same error as the final cyclical model, as the error is similar for λ = 0 and λ = 1. Additionally, it appears that most snapshots do not lie in the same minimum as the final model. Thus the snapshots are likely to misclassify different samples. Conversely, the first three snapshots achieve much higher error than the final model. This can be observed by the sharp minima around λ = 1, which suggests that mixing in any amount of the snapshot parameters will worsen performance. While the final two snapshots achieve low error, the figures suggests that they lie in the same minimum as the final model, and therefore likely add limited diversity to the ensemble.
Activation space. To further explore the diversity of models, we compute the pairwise correlation of softmax outputs for every pair of snapshots. Figure 6 displays the average correlation for both cyclic snapshots and non-cyclical snapshots. Firstly, there are large correlations between the last 3 snapshots of the non-cyclic training schedule (right). These snapshots are taken after dropping the learning rate, suggesting that each snapshot has converged to the same minimum. Though there is more diversity amongst the earlier snapshots, these snapshots have much higher error rates and are therefore not ideal for ensembling. Conversely, there is less correlation between all cyclic snapshots (left). Because all snapshots have similar accuracy (as can be seen in Figure 5), these differences in predictions can be exploited to create effective ensembles.
5 DISCUSSION
We introduce Snapshot Ensembling, a simple method to obtain ensembles of neural networks without any additional training cost. Our method exploits the ability of SGD to converge to and escape from local minima as the learning rate is lowered, which allows the model to visit several weight assignments that lead to increasingly accurate predictions over the course of training. We harness this power with the cyclical learning rate schedule proposed by Loshchilov & Hutter (2016), saving model snapshots at each point of convergence. We show in several experiments that all snapshots are accurate, yet produce different predictions from one another, and therefore are well suited for test-time ensembles. Ensembles of these snapshots significantly improve the state-of-the-art on CIFAR-10, CIFAR-100 and SVHN. Future work will explore combining Snapshot Ensembles with traditional ensembles. In particular, we will investigate how to balance growing an ensemble with new models (with random initializations) and refining existing models with further training cycles under a fixed training budget.
ACKNOWLEDGEMENTS
We thank Ilya Loshchilov and Frank Hutter for their insightful comments on the cyclic cosineshaped learning rate. The authors are supported in part by the, III-1618134, III-1526012, IIS1149882 grants from the National Science Foundation, US Army Research Office W911NF-141-0477, and the Bill and Melinda Gates Foundation.
A. Single model and Snapshot Ensemble performance over time
In Figures 7-9, we compare the test error of Snapshot Ensembles with the error of individual model snapshots. The blue curve shows the test error of a single model snapshot using a cyclic cosine learning rate. The green curve shows the test error when ensembling model snapshots over time. (Note that, unlike Figure 3, we construct these ensembles beginning with the earliest snapshots.) As a reference, the red dashed line in each panel represents the test error of single model trained for 300 epochs using a standard learning rate schedule. Without Snapshot Ensembles, in about half of the cases, the test error of final model using a cyclic learning rate—the right most point in the blue curve—is no better than using a standard learning rate schedule.
One can observe that under almost all settings, complete Snapshot Ensembles—the right most points of the green curves—outperform the single model baselines. In many cases, ensembles of just 2 or 3 model snapshots are able to match the performance of the single model trained with a standard learning rate. Not surprisingly, the ensembles of model snapshots consistently outperform any of its members, yielding a smooth curve of test error over time.

The work presented in this paper proposes a method to get an ensemble of neural networks at no extra training cost (i.e., at the cost of training a single network), by saving snapshots of the network during training. Network is trained using a cyclic (cosine) learning rate schedule; the snapshots are obtained when the learning rate is at the lowest points of the cycles. Using these snapshot ensembles, they show gains in performance over a single network on the image classification task on a variety of datasets.
Positives:
1. The work should be easy to adopt and re-produce, given the simple techinque and the experimental details in the paper.
2. Well written paper, with clear description of the method and thorough experiments.
Suggestions for improvement / other comments:
1. While it is fair to compare against other techniques assuming a fixed computational budget, for a clear perspective, thorough comaprisons with "true ensembles" (i.e., ensembles of networks trained independently) should be provided.
Specificially, Table 4 should be augmented with results from "true ensembles".
2. Comparison with true ensembles is only provided for DenseNet-40 on CIFAR100 in Figure 4. The proposed snapshot ensemble achieves approximately 66% of the improvement of "true ensemble" over the single baseline model. This is not reflected accurately in the authors' claim in the abstract: "[snapshot ensembles] **almost match[es]** the results of far more expensive independently trained [true ensembles]."
3. As mentioned before: to understand the diversity of snapshot ensembles, it would help to the diversity against different ensembling technique, e.g. (1) "true ensembles", (2) ensembles from dropout as described by Gal et. al, 2016 (Dropout as a Bayesian Approximation).


This work presents a general unsupervised learning method to improve the accuracy of sequence to sequence (seq2seq) models. In our method, the weights of the encoder and decoder of a seq2seq model are initialized with the pretrained weights of two language models and then fine-tuned with labeled data. We apply this method to challenging benchmarks in machine translation and abstractive summarization and find that it significantly improves the subsequent supervised models. Our main result is that the pretraining accelerates training and improves generalization of seq2seq models, achieving state-of-the-art results on the WMT English→German task, surpassing a range of methods using both phrase-based machine translation and neural machine translation. Our method achieves an improvement of 1.3 BLEU from the previous best models on both WMT’14 and WMT’15 English→German. On summarization, our method beats the supervised learning baseline.
This work presents a general unsupervised learning method to improve the accuracy of sequence to sequence (seq2seq) models. In our method, the weights of the encoder and decoder of a seq2seq model are initialized with the pretrained weights of two language models and then fine-tuned with labeled data. We apply this method to challenging benchmarks in machine translation and abstractive summarization and find that it significantly improves the subsequent supervised models. Our main result is that the pretraining accelerates training and improves generalization of seq2seq models, achieving state-of-the-art results on the WMT English→German task, surpassing a range of methods using both phrase-based machine translation and neural machine translation. Our method achieves an improvement of 1.3 BLEU from the previous best models on both WMT’14 and WMT’15 English→German. On summarization, our method beats the supervised learning baseline.
1 INTRODUCTION
Sequence to sequence (seq2seq) models (Sutskever et al., 2014; Cho et al., 2014; Kalchbrenner & Blunsom, 2013; Allen, 1987; Ñeco & Forcada, 1997) are extremely effective on a variety of tasks that require a mapping between a variable-length input sequence to a variable-length output sequence. The main weakness of sequence to sequence models, and deep networks in general, lies in the fact that they can easily overfit when the amount of supervised training data is small.
In this work, we propose a simple and effective technique for using unsupervised pretraining to improve seq2seq models. Our proposal is to initialize both encoder and decoder networks with pretrained weights of two language models. These pretrained weights are then fine-tuned with the labeled corpus.
We benchmark this method on machine translation for English→German and abstractive summarization on CNN and Daily Mail articles. Our main result is that a seq2seq model, with pretraining, exceeds the strongest possible baseline in both neural machine translation and phrase-based machine translation. Our model obtains an improvement of 1.3 BLEU from the previous best models on both WMT’14 and WMT’15 English→German. On abstractive summarization, our method achieves competitive results to the strongest baselines.
We also perform ablation study to understand the behaviors of the pretraining method. Our study confirms that among many other possible choices of using a language model in seq2seq with attention, the above proposal works best. Our study also shows that, for translation, the main gains come from the improved generalization due to the pretrained features, whereas for summarization the gains come from the improved optimization due to pretraining the encoder which has been unrolled for hundreds of timesteps. On both tasks, our proposed method always improves generalization on the test sets.
∗Work done as an intern on Google Brain.
2 UNSUPERVISED PRETRAINING FOR SEQUENCE TO SEQUENCE LEARNING
In the following section, we will describe our basic unsupervised pretraining procedure for sequence to sequence learning and how to modify sequence to sequence learning to effectively make use of the pretrained weights. We then show several extensions to improve the basic model.
2.1 BASIC PROCEDURE
Given an input sequence x1, x2, ..., xm and an output sequence yn, yn−1, ..., y1, the objective of sequence to sequence learning is to maximize the likelihood p(yn, yn−1, ..., y1|x1, x2, ..., xm). Common sequence to sequence learning methods decompose this objective as p(yn, yn−1, ..., y1|x1, x2, ..., xm) = ∏n t=1 p(yt|yt−1, ..., y1;x1, x2, ..., xm).
In sequence to sequence learning, an RNN encoder is used to represent x1, ..., xm as a hidden vector, which is given to an RNN decoder to produce the output sequence. Our method is based on the observation that without the encoder, the decoder essentially acts like a language model on y’s. Similarly, the encoder with an additional output layer also acts like a language model. Thus it is natural to use trained languages models to initialize the encoder and decoder.
Therefore, the basic procedure of our approach is to pretrain both the seq2seq encoder and decoder networks with language models, which can be trained on large amounts of unlabeled text data. This can be seen in Figure 1, where the parameters in the shaded boxes are pretrained. In the following we will describe the method in detail using machine translation as an example application.
First, two monolingual datasets are collected, one for the source side language, and one for the target side language. A language model (LM) is trained on each dataset independently, giving an LM trained on the source side corpus and an LM trained on the target side corpus.
After two language models are trained, a multi-layer seq2seq model M is constructed. The embedding and first LSTM layers of the encoder and decoder are initialized with the pretrained weights. To be even more efficient, the softmax of the decoder is initialized with the softmax of the pretrained target side LM.
2.2 IMPROVING THE MODEL
We also employ three additional methods to further improve the model above. The three methods are: a) Monolingual language modeling losses, b) Residual connections and c) Attention over multiple layers (see Figure 2).
Monolingual language modeling losses: After the seq2seq model M is initialized with the two LMs, it is fine-tuned with a labeled dataset. To ensure that the model does not overfit the labeled data, we regularize the parameters that were pretrained by continuing to train with the monolingual language modeling losses. The seq2seq and language modeling losses are weighted equally.
Residual connections: As described, the input vector to the decoder softmax layer is a random vector because the high level (non-first) layers of the LSTM are randomly initialized. This slows down training and introduces random gradients to the pretrained parameters, reducing the effectiveness of pretraining. To circumvent this issue, we use a residual connection from the output of the first LSTM layer directly to the input of the softmax (see Figure 2-a).
Attention over multiple layers: In all our models, we use an attention mechanism (Bahdanau et al., 2015), where the model attends over both top and first layer (see Figure 2-b). More concretely, given a query vector qt from the decoder, encoder states from the first layer h11, . . . , h 1 T , and encoder states from the last layer hL1 , . . . , h L T , we compute the attention context vector ct as follows:
αi = exp(qt · hNi )∑T j=1 exp(qt · hNj )
c1t = T∑ i=1 αih 1 i c N t = T∑ i=1 αih N i ct = [c 1 t ; c N t ]
Note that attention weights αi are only computed once using the top level encoder states.
We also experiment with passing the attention vector ct as input into the next timestep (Luong et al., 2015b). Instead of passing c into the first LSTM layer, we pass it as input to the second LSTM layer by concatenating it with the output of the first LSTM layer.
We use all three improvements in our experiments. However, in general we notice that the benefits of the attention modifications are minor in comparison with the benefits of the additional language modeling objectives and residual connections.
3 EXPERIMENTS
In the following section, we apply our approach to two important tasks in seq2seq learning: machine translation and abstractive summarization. On each task, we compare against the previous best systems. We also perform ablation experiments to understand the behavior of each component of our method.
3.1 MACHINE TRANSLATION
Dataset and Evaluation: For machine translation, we evaluate our method on the WMT English→German task (Bojar et al., 2015). We used the WMT 14 training dataset, which is slightly smaller than the WMT 15 dataset. Because the dataset has some noisy examples, we used a language detection system to filter the training examples. Sentences pairs where either the source was not English or the target was not German were thrown away. This resulted in around 4 million training examples. Following Sennrich et al. (2015b), we use subword units (Sennrich et al., 2015a)
with 89500 merge operations, giving a vocabulary size around 90000. The validation set is the concatenated newstest2012 and newstest2013, and our test sets are newstest2014 and newstest2015. Evaluation on the validation set was with case-sensitive BLEU (Papineni et al., 2002) on tokenized text using multi-bleu.perl. Evaluation on the test sets was with case-sensitive BLEU on detokenized text using mteval-v13a.pl. The monolingual training datasets are the News Crawl English and German corpora, each of which has more than a billion tokens.
Experimental settings: The language models were trained in the same fashion as (Jozefowicz et al., 2016) We used a 1 layer 4096 dimensional LSTM with the hidden state projected down to 1024 units (Sak et al., 2014) and trained for one week on 32 Tesla K40 GPUs. Our seq2seq model was a 3 layer model, where the second and third layers each have 1000 hidden units. The monolingual objectives, residual connection, and the modified attention were all used. We used the Adam optimizer (Kingma & Ba, 2015) and train with asynchronous SGD on 16 GPUs for speed. We used a learning rate of 5e-5 which is multiplied by 0.8 every 50K steps after an initial 400K steps, gradient clipping with norm 5.0 (Pascanu et al., 2013), and dropout of 0.2 on non-recurrent connections (Zaremba et al., 2014). We used early stopping on validation set perplexity. A beam size of 10 was used for decoding. Our ensemble is constructed with the 5 best performing models on the validation set, which are trained with different hyperparameters.
Results: Table 1 shows the results of our method in comparison with other baselines. Our method achieves a new state-of-the-art for single model performance on both newstest2014 and newstest2015, significantly outperforming the competitive semi-supervised backtranslation technique (Sennrich et al., 2015b). Equally impressive is the fact that our best single model outperforms the previous state of the art ensemble of 4 models. Our ensemble of 5 models matches or exceeds the previous best ensemble of 12 models.
Ablation study: In order to better understand the effects of pretraining, we conducted an ablation study by modifying the pretraining scheme. Figure 3 shows the drop in validation BLEU of various ablations compared with the full model. The full model uses LMs trained with monolingual data to initialize the encoder and decoder, in addition to the language modeling objective. In the following, we interpret the findings of the study. Note that some findings are specific to the translation task.
Given the results from the ablation study, we can make the following observations:
• Pretraining the decoder is better than pretraining the encoder: Only pretraining the encoder leads to a 1.6 BLEU point drop while only pretraining the decoder leads to a 1.0 BLEU point drop. • Pretrain as much as possible because the benefits compound: given the drops of no pre-
training at all (−2.0) and only pretraining the encoder (−1.6), the additive estimate of the drop of only pretraining the decoder side is −2.0 − (−1.6) = −0.4; however the actual drop is −1.0 which is a much larger drop than the additive estimate. • Pretraining the softmax is important: Pretraining only the embeddings and first LSTM layer
gives a large drop of 1.6 BLEU points.
• The language modeling objective is a strong regularizer: The drop in BLEU points of pretraining the entire model and not using the LM objective is as bad as using the LM objective without pretraining.
• Pretraining on a lot of unlabeled data is essential for learning to extract powerful features: If the model is initialized with LMs that are pretrained on the source part and target part of the parallel corpus, the drop in performance is as large as not pretraining at all. However, performance remains strong when pretrained on the large, non-news Wikipedia corpus.
To understand the contributions of unsupervised pretraining vs. supervised training, we track the performance of pretraining as a function of dataset size. For this, we trained a a model with and without pretraining on random subsets of the English→German corpus. Both models use the additional LM objective. The results are summarized in Figure 4. When a 100% of the labeled data is used, the gap between the pretrained and no pretrain model is 2.0 BLEU points. However, that gap grows when less data is available. When trained on 20% of the labeled data, the gap becomes 3.8 BLEU points. This demonstrates that the pretrained models degrade less as the labeled dataset becomes smaller.
3.2 ABSTRACTIVE SUMMARIZATION
Dataset and Evaluation: For a low-resource abstractive summarization task, we use the CNN/Daily Mail corpus from (Hermann et al., 2015). Following Nallapati et al. (2016), we modify the data collection scripts to restore the bullet point summaries. The task is to predict the bullet point summaries from a news article. The dataset has fewer than 300K document-summary pairs. To compare against Nallapati et al. (2016), we used the anonymized corpus. However, for our ablation study, we used the non-anonymized corpus.1 We evaluate our system using full length ROUGE (Lin, 2004). For the anonymized corpus in particular, we considered each highlight as a separate sentence following Nallapati et al. (2016). In this setting, we used the English Gigaword corpus (Napoles et al., 2012) as our larger, unlabeled “monolingual” corpus, although all data used in this task is in English.
Experimental settings: We use subword units (Sennrich et al., 2015a) with 31500 merges, resulting in a vocabulary size of about 32000. We use up to the first 600 tokens of the document and predict the entire summary. Only one language model is trained and it is used to initialize both the encoder and decoder, since the source and target languages are the same. However, the encoder and decoder are not tied. The LM is a one-layer LSTM of size 1024 trained in a similar fashion to Jozefowicz et al. (2016). For the seq2seq model, we use the same settings as the machine translation experiments. The only differences are that we use a 2 layer model with the second layer having 1024 hidden units, and that the learning rate is multiplied by 0.8 every 30K steps after an initial 100K steps.
Results: Table 2 summarizes our results on the anonymized version of the corpus. Our pretrained model is only able to match the previous baseline seq2seq of Nallapati et al. (2016). However, our model is a unidirectional LSTM while they use a bidirectional LSTM. They also use a longer context of 800 tokens, whereas we used a context of 600 tokens due to GPU memory issues. Furthermore, they use pretrained word2vec (Mikolov et al., 2013) vectors to initialize their word embeddings. As we show in our ablation study, just pretraining the embeddings itself gives a large improvement.
Ablation study: We performed an ablation study similar to the one performed on the machine translation model. The results are reported in Figure 5. Here we report the drops on ROUGE-1, ROUGE-2, and ROUGE-L on the non-anonymized validation set.
Given the results from our ablation study, we can make the following observations:
• Pretraining improves optimization: in contrast with the machine translation model, it is more beneficial to only pretrain the encoder than only the decoder of the summarization model. One interpretation is that pretraining enables the gradient to flow much further back in time than randomly initialized weights. This may also explain why pretraining on the parallel corpus is no worse than pretraining on a larger monolingual corpus.
• The language modeling objective is a strong regularizer: A model without the LM objective has a significant drop in ROUGE scores.
Human evaluation: As ROUGE may not be able to capture the quality of summarization, we also performed a small qualitative study to understand the human impression of the summaries produced by different models. We took 200 random documents and compared the performance of
1We encourage future researchers to use the non-anonymized version because it is a more realistic summarization setting with a larger vocabulary. Our numbers on the non-anonymized test set are 35.56 ROUGE-1, 14.60 ROUGE-2, and 25.08 ROUGE-L. We did not consider highlights as separate sentences.
a pretrained and non-pretrained system. The document, gold summary, and the two system outputs were presented to a human evaluator who was asked to rate each system output on a scale of 1-5 with 5 being the best score. The system outputs were presented in random order and the evaluator did not know the identity of either output. The evaluator noted if there were repetitive phrases or sentences in either system outputs. Unwanted repetition was also noticed by Nallapati et al. (2016).
Table 3 and 4 show the results of the study. In both cases, the pretrained system outperforms the system without pretraining in a statistically significant manner. The better optimization enabled by pretraining improves the generated summaries and decreases unwanted repetition in the output.
4 RELATED WORK
Unsupervised pretraining has been intensively studied in the past years, most notably is the work by Dahl et al. (2012) who found that pretraining with deep belief networks improved feedforward acoustic models. More recent acoustic models have found pretraining unnecessary (Xiong et al.,
2016; Zhang et al., 2016; Chan et al., 2015), probably because the reconstruction objective of deep belief networks is too easy. In contrast, we find that pretraining language models by next step prediction significantly improves seq2seq on challenging real world datasets.
Despite its appeal, unsupervised learning is rarely shown to improve supervised training. Dai & Le (2015) was amongst the rare studies which showed the benefits of pretraining in a semi-supervised learning setting. Their method is similar to our method except that they did not have a decoder network and thus could not apply to seq2seq learning. Similarly, Zhang & Zong (2016) found it useful to add an additional task of sentence reordering of source-side monolingual data for neural machine translation. Various forms of transfer or multitask learning with seq2seq framework also have the flavors of our algorithm (Zoph et al., 2016; Luong et al., 2015a; Firat et al., 2016).
Perhaps most closely related to our method is the work by Gulcehre et al. (2015), who combined a language model with an already trained seq2seq model by fine-tuning additional deep output layers. Empirically, their method produces small improvements over the supervised baseline. We suspect that their method does not produce significant gains because (i) the models are trained independently of each other and are not fine-tuned (ii) the LM is combined with the seq2seq model after the last layer, wasting the benefit of the low level LM features, and (iii) only using the LM on the decoder side. Venugopalan et al. (2016) addressed (i) but still experienced minor improvements. Using pretrained GloVe embedding vectors (Pennington et al., 2014) had more impact.
Related to our approach in principle is the work by Chen et al. (2016) who proposed a two-term, theoretically motivated unsupervised objective for unpaired input-output samples. Though they did not apply their method to seq2seq learning, their framework can be modified to do so. In that case, the first term pushes the output to be highly probable under some scoring model, and the second term ensures that the output depends on the input. In the seq2seq setting, we interpret the first term as a pretrained language model scoring the output sequence. In our work, we fold the pretrained language model into the decoder. We believe that using the pretrained language model only for scoring is less efficient that using all the pretrained weights. Our use of labeled examples satisfies the second term. These connections provide a theoretical grounding for our work.
In our experiments, we benchmark our method on machine translation, where other unsupervised methods are shown to give promising results (Sennrich et al., 2015b; Cheng et al., 2016). In backtranslation (Sennrich et al., 2015b), the trained model is used to decode unlabeled data to yield extra labeled data. One can argue that this method may not have a natural analogue to other tasks such as summarization. We note that their technique is complementary to ours, and may lead to additional gains in machine translation. The method of using autoencoders in Cheng et al. (2016) is promising, though it can be argued that autoencoding is an easy objective and language modeling may force the unsupervised models to learn better features.
5 CONCLUSION
We presented a novel unsupervised pretraining method to improve sequence to sequence learning. The method can aid in both generalization and optimization. Our scheme involves pretraining two language models in the source and target domain, and initializing the embeddings, first LSTM layers, and softmax of a sequence to sequence model with the weights of the language models. Using our method, we achieved state-of-the-art machine translation results on both WMT’14 and WMT’15 English to German.
A key advantage of this technique is that it is flexible and can be applied to a large variety of tasks, such as summarization, where it surpasses the supervised learning baseline.
ACKNOWLEDGMENTS
We thank George Dahl, Andrew Dai, Laurent Dinh, Stephan Gouws, Geoffrey Hinton, Rafal Jozefowicz, Pooya Khorrami, Phillip Louis, Ramesh Nallapati, Arvind Neelakantan, Xin Pan, Abi See, Rico Sennrich, Luke Vilnis, Yuan Yu and the Google Brain team for their help with the project.
Source Document
SELECTED SUMMARIZATION OUTPUTS
APPENDIX
Source Document
Source Document
Source Document
Source
SELECTED ENGLISH→GERMAN OUTPUTS
Source
Source
Source

strengths:
A method is proposed in this paper to initialize the encoder and decoder of the seq2seq model using the trained weights of language models with no parallel data. After such pretraining, all weights are jointly fine-tuned with parallel labeled data with an additional language modeling loss.
It is shown that pretraining accelerates training and improves generalization of seq2seq models.
The main value of the proposed method is to leverage separate source and target corpora, contrasting the common methods of using large amounts of parallel training corpora.
weaknesses:
The objective function shown in the middle of pg 3 is highly empirical, not directly linked to how non-parallel data helps to improve the final prediction results. The paper should compare with and discuss the objective function based on expectation of cross entropy which is directly linked to improving prediction results as proposed in arXiv:1606.04646, Chen et al.:  Unsupervised Learning of Predictors from Unpaired Input-Output Samples, 2016.
The pre-training procedure proposed in this paper is also closely connected with the DNN pretraining method presented in Dahl et al. 2011, 2012. Comparisons should be made in the paper, highlighting why the proposed one is conceptually superior if the authors believe so.