Grounded Compositional Generalization with Environment Interactions

1 INTRODUCTION . Compositional generalization is a key skill for flexible and efficient learning . Humans leverage compositionality to create and recognize new combinations of familiar concepts ( Chomsky , 1957 ; Minsky , 1986 ) . Though there are many progresses for machine learning and deep learning in various areas recently ( LeCun et al. , 2015 ) , current main learning algorithms are not able to perform compositional generalization , and require many samples to train models . Such efficient learning is even more important when machines interact with the environment for grounding , because interactions are usually slow . Machine learning has been mostly developed with an assumption that training and test distributions are identical . Compositional generalization , however , is a kind of out-of-distribution generalization ( Bengio , 2017 ) , where training and test distributions are different . During training , dataset does not contain the information of the difference , so it can only be given as prior knowledge . In compositional generalization , a sample is a combination of several components . Test distribution changes as test samples are new combinations of seen components in training . For example , if we can find “ large apple ” and “ small orange ” in some environments , then we can also find “ large orange ” among multiple objects in a new environment . The recombination is enabled when an output component depends only on the corresponding input components , and invariant of other components ( please see Section 4.1 for more details ) . So there are two aspects to consider . What are the components in output , and how to find the corresponding input signals . We propose to use interactions between agent and the environment to define output components . This is analogues to model-free reinforcement learning ( Sutton & Barto , 2018 ) , where an agent does not have an environment model , but leans to act at each step during the interactions with the environment . Then we use entropy regularization ( Li et al. , 2019 ; Li & Eisner , 2019 ) to learn the minimal input components for outputs . We evaluate the approach with gSCAN dataset ( Ruis et al. , 2020 ) , which is designed to study compositional generalization in grounded agent instruction learning . Please see Figure 1 for examples . The results show the proposed approach significantly outperforms baselines in most tasks , with more than 25 % absolute average accuracy increase , and the high accuracy indicates that the proposed approach addresses the designed grounded compositional generalization problems in these tasks . We also look into the impact of entropy regularization and other changes with ablation study . We hope this work will be helpful in advancing grounded compositional generalization and artificial intelligence research . The contributions of this paper can be summarized as follows . • This is the first work to enable accurate compositional generalization in grounded instruction learning problem , serving for analyzing and understanding the mechanism . • The novelty of this paper is to find that the combination of environment interaction and entropy regularization helps the generalization . 2 RELATED WORK . Compositional generalization research has a long history , and recently there are increasing focus on this area . SCAN dataset ( Lake , 2019 ) was proposed to study compositional generalization in instruction learning . It maps a command sentence to a sequence of actions . This dataset has a property that input words and output actions have direct correspondence . Though some NLP tasks , such as machine translation , have similar property , not all problems fit to the setting . Also this dataset does not contain an environment for an agent to take actions . SCAN dataset inspired multiple approaches ( Russin et al. , 2019 ; Andreas , 2019 ; Li et al. , 2019 ; Lake , 2019 ; Gordon et al. , 2020 ; Liu et al. , 2020 ) . Some of them lead to general techniques for compositional generalization . For example , entropy regularization ( Li et al. , 2019 ) is proposed to avoid redundant dependency on input , and it is a core idea of the approach in this paper . Compositional generalization has applications in various fields such as question answering ( Andreas et al. , 2016 ; Hudson & Manning , 2019 ; Keysers et al. , 2020 ) , counting ( Rodriguez & Wiles , 1998 ; Weiss et al. , 2018 ) , systematic behaviour ( Wong & Wang , 2007 ; Brakel & Frank , 2009 ) , and hierarchical structure ( Linzen et al. , 2016 ) . Another related work is independent disentangled representation ( Higgins et al. , 2017 ; Locatello et al. , 2019 ) , but they do not address compositional generalization . Compositionality is also helpful for reasoning ( Talmor et al. , 2020 ) and continual learning ( Jin et al. , 2020 ; Li et al. , 2020 ) . Grounded SCAN ( gSCAN ) dataset was proposed to introduce environment and grounding to agent instruction learning with compositional generalization ( Ruis et al. , 2020 ) . It has a command sentence as input and a sequence of actions as output . However , the input command does not tell the specific way to act , but agent needs to understand the environment and take corresponding actions . This also avoids direct mapping between input words and output actions . Different approaches have been proposed to address this problem . As compositional generalization requires prior knowledge for distribution change , these approaches correspond to different ways to provide the prior knowledge . Andreas ( 2019 ) uses linguistic knowledge to augment training data . Kuo et al . ( 2020 ) uses external syntactic parser and WordNet . In this paper , we apply the prior knowledge for the interactions of agent and the environment . 3 PROBLEM DESCRIPTION . gSCAN dataset contains episodes of commands and actions , and each episode has an input and an output . The input includes a command and an environment . A command is a sequence of words . An environment contains an agent and a set of objects . An agent has its position and initial direction . An object has its position and attributes of color , shape and size . The output contains a sequence of actions ( Figure 1 ) . We hope to extract prior knowledge from environment interactions . We first notice that an agent should know whether it is going to change position . This means we can break an episode to a sequence of steps , where each step corresponds to an action for position change . An agent should also know the change of directions between steps . So we separate the direction change and the manner of action in a step . Therefore , we have three output for each step : direction , action and manner . We further convert the environment to be agent centered , and rotate the environment to make agent facing forward . We also assume automatic collision prevention . If the agent tries to push or pull an object to collide with other objects or boarder , then it stops . This makes us focus on addressing grounded compositional generalization problem . In summary , when we consider agent interaction with the environment , we can convert the problem to a set of step-wise label prediction problems with multiple outputs . Input contains command , environment , state , and output contains direction , action and manner . 4 APPROACH . In this section , we describe the approach for grounded compositional generalization . As we use entropy regularization in different modules , we first introduce it , then move to the model architecture . 4.1 ENTROPY REGULARIZATION . The difficulty of compositional generalization is that there might be incorrect dependency between input and output components . For example , when “ red ” and “ square ” do not appear together in any training sample , a model might learn that square is not red . However , this causes errors for compositional generalization in test . To avoid such case , we hope the representation of shape not influenced by input information of color . Entropy regularization ( Li et al. , 2019 ; Li & Eisner , 2019 ) aims at reducing entropy of a representation to avoid dependency on redundant input components . Given an representation x , we compute the L2 norm and add normal noise to each element of the representation . This decreases the channel capacity , so that the entropy for the representation reduces . We then feed the noised representation to the next layer , and add the norm to loss function . L = Loriginal + λL2 ( x ) EntReg ( x ) = x+ αN ( 0 , I ) where α is a weight of noise , positive for training and zero for inference . A representation can be fed to multiple networks , requiring different regularization for each input node . So we design entropy regularization layer , where we achieve non-linear mapping by expanding each node xi to a vector hi ∈ RH with ReLU activation , and maps it back to a node yi with linear activation , then apply entropy regularization on y to get y′ . The input x and output ERL ( x ) = y′ have the same size . hi = ffAi ( xi , H ) , yi = ff B i ( hi , 1 ) , y ′ = EntReg ( y ) . We write ff ( x , K ) for feed-forward neural network with x as input and K as output size . We use EntReg for word embeddings and ERL for environment inputs . 4.2 MODEL ARCHITECTURE . Input contains one-hot representation of a command xc with sentence length n and vocabulary size V c , and a sequence of m objects . We denote object attributes xa with vector size V a , and object positions xp with vector size V p. We also has a binary state s ∈ { 0 , 1 } indicating whether it is the first step of an episode . Output has three types . Direction yd has Cd classes , action ya has Ca classes , and manner ym has Cm classes . xc = xc1 , . . . , x c n ∈ RV c×n , xa = xa1 , . . . , x a m ∈ RV a×m , xp = xp1 , . . . , x p m ∈ RV p×m . The model has the following modules , as summarized in Figure 2 and Algorithm 1 . Command module ( CM ) This module takes a command xc as input and returns a representation r with size K , r = CM ( xc , K ) . r is expected to be the embedding of certain type of keyword , e.g . action , color , etc . To enable compositional generalization , we hope to separate representations of types ( e.g . color ) and values ( e.g . red ) . We use types for attention maps and values for attended values in attention mechanism . For a new combination of values , the types are still recognizable , so that attention maps are correct and the corresponding values are extracted ( Li et al. , 2019 ) . We design that each word xci in command x c has two embeddings for type ti and value vi with corresponding embedding matrices , Et ∈ Rkt×V c , Ev ∈ Rkv×V c , where kt and kv are embedding sizes , respectively . We then apply entropy regularization on both t and v to reduce redundant dependency . t = EntReg ( Etxc ) ∈ Rkt×n , v = EntReg ( Evxc ) ∈ Rkv×n . We use attention mechanism with a query q ∈ Rkt as learnable parameters , keys t , values v , and temperature τ c ∈ R. We compute score wc and attention map ac . The attended value u is fed to a feed-forward network with output size K. wc = qt ∈ Rn , ac = Softmax ( wc/τ c ) ∈ Rn , u = vac ∈ Rkv , r = ffθ ( u , K ) ∈ RK . Grounding module ( GM ) This module finds target in the environment according to the command . We use N command modules for queries . For each query ri , we have an dedicated entropy regularization layer on attributes x′ai , because different queries correspond to different type of attributes , and other attributes would be redundant . We then compute a score wgi . For i = 1 , . . . , N , ri = CM ( xc , V a ) ∈ RV a x′ai = ERL ( x a ) ∈ RV a×m , wgi = rix ′a i ∈ Rm . Algorithm 1 The proposed approach . Input includes command xc , state s , object attributes xa and object positions xp . K is an embedding size . Please see Section 4 for more information . Command module Input : xc , K Output : r : embedding 1 : t = EntReg ( Etxc ) 2 : v = EntReg ( Evxc ) 3 : wc = qt 4 : ac = Softmax ( wc/τ c ) 5 : u = vac 6 : r = ffθ ( u , K ) Grounding module Input : xc , xa , xp Output : p : target position 1 : for i = 1 , . . . , N do 2 : ri = CM ( xc , V a ) 3 : x′ai = ERL ( x a ) 4 : wgi = rix ′a i 5 : end for 6 : wg = ∑N i=1 w g i 7 : ag = Softmax ( wg/τg ) 8 : p = xpag Prediction module Input : xc , xa , xp , s , K Output : ŷ : prediction 1 : p = GM ( xc , xa , xp ) 2 : for each output type i do 3 : ki = CM ( xc , K ) 4 : for each node j do 5 : p′i , j = ERL ( p ) 6 : zi , j = [ kTi , p ′T i , j , s T ] T 7 : li , j = ff φ i , j ( zi , j , 1 ) 8 : end for 9 : ŷi = Softmax ( li ) 10 : end for The scores are added as wg . Attention map ag is computed by Softmax with temperature τg ∈ R , and is applied to get the attended object position p. wg = N∑ i=1 wgi ∈ R m , ag = Softmax ( wg/τg ) ∈ Rm , p = xpag ∈ RV p . Prediction module ( PM ) Prediction module takes command , environment and state as input , and outputs a prediction . We have three separate prediction modules for direction , action and manner , respectively . Modules correspond to different keywords but share the same grounded target and state . So for each prediction module , we use one command module to extract a keyword k with size K for the prediction . We also have a environment module for target object position p , and all the prediction modules share it as an input . There is a input of state s , and this is also a shared input for each prediction module . We build a dedicated feed-forward neural network from input to each output node without weight sharing . In each network of output type i and node j , we use entropy regularization layer for target position p. This is because different output nodes may need to be computed from different input components , and other components are redundant , so we hope to reduce dependency of each output node to input nodes ( more discussion in Section 6.1 ) . Then all the inputs are concatenated to form a vector z with size L = K + V p + 1. ki = CM ( xc , K ) ∈ RK , p′i , j = ERL ( p ) ∈ RV p , zi , j = [ k T i , p ′T i , j , s T ] T ∈ RL . We then feed it to another feed-forward network to get a output node li , j . They are concatenated to form a logit li , and we use Softmax to output ŷi . Ci is the number of classes for the output type . li , j = ff φ i , j ( zi , j , 1 ) ∈ R , li = [ li,1 , . . . , li , Ci ] ∈ R Ci , ŷi = Softmax ( li ) ∈ RCi We use cross entropy and the norms for entropy regularization with weight λ as training objective .

Summary: This paper tries to address a very important problem, compositional generalization in grounded agent instruction learning. It proposes to use interactions between agent and the environment to define output components, and entropy regularization to reduce redundant dependency on input. It shows significant improvements in most of gSCAN tasks. The paper also has an ablation study that investigates the effectiveness of entropy regularization and other factors.