Image captioning with weakly-supervised attention penalty
Techniques for training a machine-learning (ML) model for captioning images are disclosed. A plurality of feature vectors and a plurality of visual attention maps are generated by a visual model of the ML model based on an input image. Each of the plurality of feature vectors correspond to different regions of the input image. A plurality of caption attention maps are generated by an attention model of the ML model based on the plurality of feature vectors. An attention penalty is calculated based on a comparison between the caption attention maps and the visual attention maps. A loss function is calculated based on the attention penalty. One or both of the visual model and the attention model are trained using the loss function.
1. A method of training a machine-learning (ML) model for captioning images, the method comprising:
generating, by a visual model of the ML model and based on an input image, a plurality of feature vectors and a plurality of visual attention maps, wherein each of the plurality of feature vectors corresponds to different regions of the input image;
generating, by an attention model of the ML model, a plurality of caption attention maps based on the plurality of feature vectors;
calculating an attention penalty based on a comparison between the caption attention maps and the visual attention maps;
calculating a loss function based on the attention penalty; and
training one or both of the visual model and the attention model using the loss function;
wherein the training one or both of the visual model and the attention model using the loss function includes modifying a set of weights associated with the visual model and/or the attention model by an amount proportional to a magnitude of the loss function.
2. The method of claim 1 , further comprising:
generating, by a language model of the ML model, a predicted caption based on the plurality of feature vectors.
3. The method of claim 1 , wherein the plurality of visual attention maps and the plurality of caption attention maps are probability distributions.
4. The method of claim 1 , wherein the plurality of visual attention maps are generated by a gradients-based weakly supervised model.
5. The method of claim 1 , wherein the plurality of visual attention maps are generated by a class activation mapping algorithm.
6. The method of claim 2 , further comprising:
calculating a caption loss based on a comparison between the predicted caption and a caption label provided by a user, wherein the loss function is calculated based on at least the caption loss and the attention penalty.
7. The method of claim 2 , further comprising:
training the language model using the loss function.
8. The method of claim 2 , wherein each visual attention map of the plurality of visual attention maps corresponds to a parsed visual word obtained from a caption using a category mapping dictionary.
9. The method of claim 2 , further comprising:
determining a word of a caption is not a noun or in an object category dictionary; and
setting the attention penalty to zero.
10. A non-transitory computer-readable medium comprising instructions that, when executed by one or more processors, cause the one or more processors to perform operations comprising:
generating, by a visual model of a machine-learning (ML) model for captioning images and based on an input image, a plurality of feature vectors and a plurality of visual attention maps, wherein each of the plurality of feature vectors corresponds to different regions of the input image;
generating, by an attention model of the ML model, a plurality of caption attention maps based on the plurality of feature vectors;
calculating an attention penalty based on a comparison between the caption attention maps and the visual attention maps;
calculating a loss function based on the attention penalty; and
training one or both of the visual model and the attention model using the loss function;
wherein the training one or both of the visual model and the attention model using the loss function includes modifying a set of weights associated with the visual model and/or the attention model by an amount proportional to a magnitude of the loss function.
11. The non-transitory computer-readable medium of claim 10 , wherein the operations further comprise:
generating, by a language model of the ML model, a predicted caption based on the plurality of feature vectors.
12. The non-transitory computer-readable medium of claim 10 , wherein the plurality of visual attention maps and the plurality of caption attention maps are probability distributions.
13. The non-transitory computer-readable medium of claim 11 , wherein the operations further comprise:
calculating a caption loss based on a comparison between the predicted caption and a caption label provided by a user, wherein the loss function is calculated based at least on the caption loss and the attention penalty.
14. The non-transitory computer-readable medium of claim 11 , wherein the operations further comprise:
training the language model using the loss function.
15. A system comprising:
one or more processors; and
a computer-readable medium comprising instructions that, when executed by the one or more processors, cause the one or more processors to perform operations comprising:
generating, by a visual model of a machine-learning (ML) model for captioning images and based on an input image, a plurality of feature vectors and a plurality of visual attention maps, wherein each of the plurality of feature vectors corresponds to different regions of the input image;
generating, by an attention model of the ML model, a plurality of caption attention maps based on the plurality of feature vectors;
calculating an attention penalty based on a comparison between the caption attention maps and the visual attention maps;
calculating a loss function based on the attention penalty; and
training one or both of the visual model and the attention model using the loss function;
wherein the training one or both of the visual model and the attention model using the loss function includes modifying a set of weights associated with the visual model and/or the attention model by an amount proportional to a magnitude of the loss function.
16. The system of claim 15 , wherein the operations further comprise:
generating, by a language model of the ML model, a predicted caption based on the plurality of feature vectors.
17. The system of claim 15 , wherein the visual model is a pre-trained convolutional neural network fine-tuned using labels generating by parsing captions and mapping visual concepts to object categories.
18. The system of claim 15 , wherein the attention penalty comprises a sum of a difference map determined between the caption attention maps and the visual attention maps.
19. The system of claim 16 , wherein the operations further comprise:
calculating a caption loss based on a comparison between the predicted caption and a caption label provided by a user, wherein the loss function is calculated based on at least the caption loss and the attention penalty.
20. The system of claim 16 , wherein the operations further comprise:
training the language model using the loss function.