Learning Compact Reward for Image Captioning

Adversarial learning has shown its advances in generating natural and diverse descriptions in image captioning . However , the learned reward of existing adversarial methods is vague and ill-defined due to the reward ambiguity problem . In this paper , we propose a refined Adversarial Inverse Reinforcement Learning ( rAIRL ) method to handle the reward ambiguity problem by disentangling reward for each word in a sentence , as well as achieve stable adversarial training by refining the loss function to shift the stationary point towards Nash equilibrium . In addition , we introduce a conditional term in the loss function to mitigate mode collapse and to increase the diversity of the generated descriptions . Our experiments on MS COCO show that our method can learn compact reward for image captioning . 1 INTRODUCTION . Image captioning is a task of generating descriptions of a given image in natural language . In a general encoder-decoder structure ( Vinyals et al. , 2015 ) , image features are encoded in a CNN and decoded into a caption in a word by word manner . Based on the loss function , standard approaches to the problem could be divided into three categories : MLE ( Maximum Likelihood Estimation ) , RL ( Reinforcement Learning ) and GAN ( Generative Adversarial Network ) . Early proposed methods were based on MLE function and made improvements by designing specific model structure ( Xu et al. , 2015 ) . MLE adopts the cross-entropy loss and learns a one-hot distribution for each word in the sentence . By maximizing the probability of the ground truth word whilst suppressing other reasonable vocabularies , the probability distribution learned by MLE tends to be sparse and the generated captions have limited diversity ( Dai et al. , 2017 ) . On the other hand , RL has advantages in boosting the model performance by optimizing the handcrafted metrics ( Rennie et al. , 2017 ; Liu et al. , 2017 ; Chen et al. , 2019 ) . However , due to the reward hacking problem , RL maximizes the reward in an unintended way and fails to produce human-like descriptions ( Li et al. , 2019a ) . Considering naturalness and diversity of the generated captions , GAN has raised attention in image captioning for its capability of producing descriptions that are indistinguishable from human-written ones ( Dai et al. , 2017 ; Shetty et al. , 2017 ; Chen et al. , 2019 ; Dognin et al. , 2019 ) . In image captioning , the generator of GAN learns true data distribution by maximizing the reward function learned from a discriminator , and the discriminator distinguishes the generated sample from the true data . The adversarial training converges to an equilibrium point ( i.e. , Nash equilibrium ) at which both the generator and discriminator can not improve ( Goodfellow et al. , 2014 ) . As shown in Figure 1 , the learned distribution of GAN is closer to the ground truth distribution than that of other methods ( i.e. , MLE and RL ) on different splits . However , previous work of adversarial networks in image captioning gives one reward function D for a complete sentence consisting of n words . This strategy causes the reward ambiguity problem ( Ng et al. , 1999 ) since which word ( s ) causes the reward to increase or decrease is not accounted for , and thus there are many optimal policies that determine the sentence can explain one reward . An example is that each image in MS COCO has five ground truth captions . Although these captions may vary in formats , each caption has the same reward value in GAN . From the perspective on the system level , learning sentence-level reward from different image-caption pairs is analogous to learning reward of a trajectory from different system dynamics , which makes the discriminator unable to distinguish the true reward functions from those shaped by the environment dynamics ( Fu et al. , 2018 ) . Facing the challenge , we adopt AIRL ( Fu et al. , 2018 ) to solve the reward ambiguity problem by disentangling reward for each action ( i.e. , word in a sentence ) from different image-caption pairs and learning a compact reward function . compact means words with similar semantics , such as children and kids , correspond to close reward values . Driven by the compact reward function of the discriminator , the generator learns the optimal policy and thus produces qualitative descriptions . However , there are still two major problems to address : 1 ) AIRL is difficult to converge to Nash equilibrium using policy gradient ( See Section 4.2 for details ) ; 2 ) AIRL is designed without mode control , and thus the outputs have limited diversity , which is a commonly encountered issue called mode collapse ( Mirza & Osindero , 2014 ) . In this paper , we propose a refined AIRL method to learn a compact reward function for each word , as well as achieve stable adversarial training by refining the loss function to shift the stationary point towards Nash equilibrium . The refined method makes it possible to reach the equilibrium point for a non-concave model function of the generator . In addition , a conditional term is introduced in the loss function to mitigate mode collapse and to increase the diversity of the generated descriptions . Both the caption evaluator ( i.e. , discriminator ) ( Cui et al. , 2018 ; Sharif et al. , 2018 ) and the generator are cast into this unified framework , where the discriminator evaluates captions using a learned compact reward function , and the generator produces qualitative image descriptions . We demonstrate the effectiveness of our method in the experiments . 2 RELATED WORK . Image Captioning . The development of image captioning can be summarized into two directions : model structure design ( Lu et al. , 2017 ; Yao et al. , 2018 ) and loss function construction ( Rennie et al. , 2017 ; Ren et al. , 2017 ) . In the methods based on model structure design , attention mechanism and the fusion of visual and semantic information are the key focus . Lu et al . ( 2017 ; 2018 ) proposed a sentinel gate to learn adaptive attention between visual content and non-visual text . Yao et al . ( 2018 ) explored the role of visual relationship in image captioning . On the other hand , methods based on loss function construction focus on optimization of the loss function . Rennie et al . ( 2017 ) optimized on non-differentiable evaluation metric using policy gradient , and improved scores of these metrics on various models . Ren et al . ( 2017 ) designed an embedding reward under actor-critic reinforcement learning . Similarly , we address the construction of loss functions , and thus our algorithm can be built on existing model structures . See Appendix F for a short discussion about different loss functions . Adversarial Methods for Image Captioning . Adversarial methods are known for producing plausible samples by training the generator and the discriminator in an adversarial manner ( Goodfellow et al. , 2014 ) . In image captioning , the discriminator is formed as a binary classifier that distinguishes the generated sentence from the ground truth , while the generator produce captions that can fool the discriminator . Conditional GAN was proposed in ( Dai et al. , 2017 ) to improve the naturalness and diversity of generated captions . CNN and RNN based discriminators were introduced in ( Chen et al. , 2019 ) . However , existing methods estimate a reward function for the complete sentence consisting of n words , where multiple optimal policies that determine the sentences can correspond to one reward ( Ng et al. , 1999 ) . Thus the learned reward is ambiguous and ill-defined . We solve this problem by recovering a compact reward function for each word in the sentence under a refined AIRL framework . Although AIRL has been utilized to solve problems in other fields ( Wang et al. , 2018 ; Li et al. , 2019b ; Shi et al. , 2018 ) , we are the first to make algorithmic improvements to AIRL such that Nash equilibrium can be reached even for a non-concave model function of the generator , and that diversity of the outputs can be increased . 3 ADVERSARIAL INVERSE REINFORCEMENT LEARNING . Due to the high variance estimate of a full sentence and the reward ambiguity problem , instead of learning reward for a complete sentence , we could learn reward distribution pθ ( wt , st ) for each word-state pair ( wt , st ) so that the true reward can be recovered at optimality ( Fu et al. , 2018 ) . In the following , we introduce how AIRL disentangles reward for each word-state pair ( wt , st ) . AIRL is an adversarial reward learning algorithm based on Maximum-Entropy-IRL . Finn et al . ( 2016 ) first proved that Maximum-Entropy-IRL is mathematically equivalent to GAN under a special form of the discriminator : Dθ ( wt , st ) = pθ ( wt , st ) pθ ( wt , st ) + π ( wt , st ) ( 1 ) pθ ( wt , st ) = exp { fθ ( wt , st ) } ( 2 ) where pθ ( wt , st ) is the actual probability distribution estimated by the discriminator . π ( wt , st ) is the policy distribution produced by the generator , which is also called vocabulary distribution under the context of image captioning . The goal is to estimate pθ ( wt , st ) that approximates the true data distribution ptrue ( wt , st ) , as well as to learn an optimal vocabulary distribution π that maximizes the reward . Subsequently , considering reward ambiguity problem , Fu et al . ( 2018 ) further extended the theory to AIRL by adding a reward shaping term hϕ into fθ ( wt , st ) , and thus disentangled reward from different system dynamics ( i.e. , image-caption pairs ) : fθ , ϕ ( wt , st ) = gθ ( wt , st ; st+1 ) + γhϕ ( st+1 ) − hϕ ( st ) ( 3 ) where gθ denotes the reward approximator that recovers the true reward up to a constant , and hϕ is the reward shaping term that preserves the optimal π. γ is a constant in range ( 0 , 1 ] . In the context of divergence minimization , the adversarial process can be represented as a min-max game ( Mescheder & Geiger , 2017 ) : min π max θ , ϕ Ewtruet ∼ptrue [ log ( Dθ , ϕ ( w true t , s true t ) ) ] + Ewt∼π [ log ( 1−Dθ , ϕ ( wt , st ) ) ] ( 4 ) where ptrue is the true data distribution and π is the vocabulary distribution learned by the generator . ( wtruet , s true t ) is the word-state pair of the true data . Despite of the capability of AIRL in disentangling reward for each word , it is difficult for the above AIRL algorithm to converge to Nash equilibrium and to produce diverse outputs through adversarial training ( See Section 4.2 for details ) . These issues can result in a non-optimal solution and lack of diversity of the generated descriptions . Aiming to learn the optimal compact reward as well as diverse captions , we refine the loss function to shift the stationary point towards Nash equilibrium and to mitigate mode collapse in the two-player game .

Summary: The authors propose using a recent method for adversarial inverse reinforcement learning (AIRL) for the task for generating high-quality image captions. Leveraging the GAN framework, a discriminator is trained to distinguish real captions from those produced by the generator, while the generator is optimized with policy-gradients (REINFORCE) to maximize the pseudo-reward from the discriminator. The main difference from prior work seems to be that the discriminator acts on a word-level, rather than sentence-level (as done, for instance in Dai et. al. 2017). Correspondingly, the generator policy is updated with the objective of 1-step reward maximization (more like contextual bandits), rather than with a long-term sequential decision-making objective (as done in Dai et. al. 2017). The evaluation is done using 2 data splits – standard and robust, with various metrics such as SPICE, CIDEr, BLEU, CHAIR. Diversity analysis and ablations are also performed to dissect the performance of the proposed approach.