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: This paper proposes a refined Adversarial Inverse Reinforcement Learning (rAIRL) to remedy the reward ambiguity by decoupling the reward for each word in a sentence, while the existing methods that utilize reinforcement learning to optimize evaluation score handle only sentence-level rewards. Furthermore, a conditional term is introduced in the loss function to avoid mode collapse and to increase the diversity of the generated captions. Throughout experiments on MS COCO show that the proposed method achieves state-of-the-art performance with several evaluation scores.