mac_cell.py

import collections
import numpy as np
import tensorflow as tf

import ops
from config import config

MACCellTuple = collections.namedtuple("MACCellTuple", ("control", "memory"))

'''
The MAC cell.

Recurrent cell for multi-step reasoning. Presented in https://arxiv.org/abs/1803.03067.
The cell has recurrent control and memory states that interact with the question 
and knowledge base (image) respectively.

The hidden state structure is MACCellTuple(control, memory)

At each step the cell performs by calling to three subunits: control, read and write.

1. The Control Unit computes the control state by computing attention over the question words.
The control state represents the current reasoning operation the cell performs.

2. The Read Unit retrieves information from the knowledge base, given the control and previous
memory values, by computing 2-stages attention over the knowledge base.

3. The Write Unit integrates the retrieved information to the previous hidden memory state,
given the value of the control state, to perform the current reasoning operation.
'''
class MACCell(tf.nn.rnn_cell.RNNCell):

    '''Initialize the MAC cell. 
    (Note that in the current version the cell is stateful -- 
    updating its own internals when being called) 
    
    Args:
        vecQuestions: the vector representation of the questions. 
        [batchSize, ctrlDim]

        questionWords: the question words embeddings. 
        [batchSize, questionLength, ctrlDim]

        questionCntxWords: the encoder outputs -- the "contextual" question words. 
        [batchSize, questionLength, ctrlDim]

        questionLengths: the length of each question.
        [batchSize]
        
        memoryDropout: dropout on the memory state (Tensor scalar).
        readDropout: dropout inside the read unit (Tensor scalar).
        writeDropout: dropout on the new information that gets into the write unit (Tensor scalar).
        
        batchSize: batch size (Tensor scalar).
        train: train or test mod (Tensor boolean).
        reuse: reuse cell

        knowledgeBase: 
    '''
    def __init__(self, vecQuestions, questionWords, questionCntxWords, questionLengths, 
            knowledgeBase, memoryDropout, readDropout, writeDropout, 
            batchSize, train, reuse = None):
        
        self.vecQuestions = vecQuestions
        self.questionWords = questionWords
        self.questionCntxWords = questionCntxWords
        self.questionLengths = questionLengths

        self.knowledgeBase = knowledgeBase

        self.dropouts = {}
        self.dropouts["memory"] = memoryDropout 
        self.dropouts["read"] = readDropout 
        self.dropouts["write"] = writeDropout

        self.none = tf.zeros((batchSize, 1), dtype = tf.float32)

        self.batchSize = batchSize
        self.train = train
        self.reuse = reuse

    ''' 
    Cell state size. 
    '''
    @property
    def state_size(self):
        return MACCellTuple(config.ctrlDim, config.memDim)

    '''
    Cell output size. Currently it doesn't have any outputs. 
    '''
    @property
    def output_size(self):
        return 1

    # pass encoder hidden states to control?
    '''
    The Control Unit: computes the new control state -- the reasoning operation,
    by summing up the word embeddings according to a computed attention distribution.
    
    The unit is recurrent: it receives the whole question and the previous control state,
    merge them together (resulting in the "continuous control"), and then uses that 
    to compute attentions over the question words. Finally, it combines the words 
    together according to the attention distribution to get the new control state. 
    
    Args:
        controlInput: external inputs to control unit (the question vector).
        [batchSize, ctrlDim]

        inWords: the representation of the words used to compute the attention.
        [batchSize, questionLength, ctrlDim]

        outWords: the representation of the words that are summed up. 
                  (by default inWords == outWords)
        [batchSize, questionLength, ctrlDim]

        questionLengths: the length of each question.
        [batchSize]

        control: the previous control hidden state value.
        [batchSize, ctrlDim]

        contControl: optional corresponding continuous control state
        (before casting the attention over the words).
        [batchSize, ctrlDim]

    Returns:
        the new control state
        [batchSize, ctrlDim]

        the continuous (pre-attention) control
        [batchSize, ctrlDim]
    '''
    def control(self, controlInput, inWords, outWords, questionLengths,
        control, contControl = None, name = "", reuse = None):

        with tf.variable_scope("control" + name, reuse = reuse):
            dim = config.ctrlDim

            ## Step 1: compute "continuous" control state given previous control and question.
            # control inputs: question and previous control
            newContControl = controlInput
            if config.controlFeedPrev:
                newContControl = control if config.controlFeedPrevAtt else contControl
                if config.controlFeedInputs:
                    newContControl = tf.concat([newContControl, controlInput], axis = -1)
                    dim += config.ctrlDim

                # merge inputs together
                newContControl = ops.linear(newContControl, dim, config.ctrlDim,
                    act = config.controlContAct, name = "contControl")
                dim = config.ctrlDim

            ## Step 2: compute attention distribution over words and sum them up accordingly.
            # compute interactions with question words
            interactions = tf.expand_dims(newContControl, axis = 1) * inWords
            
            # optionally concatenate words
            if config.controlConcatWords:
                interactions = tf.concat([interactions, inWords], axis = -1)
                dim += config.ctrlDim                                              

            # optional projection
            if config.controlProj:
                interactions = ops.linear(interactions, dim, config.ctrlDim, 
                    act = config.controlProjAct) 
                dim = config.ctrlDim

            # compute attention distribution over words and summarize them accordingly 
            logits = ops.inter2logits(interactions, dim)
            # self.interL = (interW, interb)

            # if config.controlCoverage:
            #     logits += coverageBias * coverage

            attention = tf.nn.softmax(ops.expMask(logits, questionLengths))
            self.attentions["question"].append(attention)

            # if config.controlCoverage:
            #     coverage += attention # Add logits instead?               

            newControl = ops.att2Smry(attention, outWords) 
            
            # ablation: use continuous control (pre-attention) instead
            if config.controlContinuous:
                newControl = newContControl

        return newControl, newContControl

    '''
    The read unit extracts relevant information from the knowledge base given the
    cell's memory and control states. It computes attention distribution over
    the knowledge base by comparing it first to the memory and then to the control.
    Finally, it uses the attention distribution to sum up the knowledge base accordingly,
    resulting in an extraction of relevant information. 

    Args:
        knowledge base: representation of the knowledge base (image). 
        [batchSize, kbSize (Height * Width), memDim]

        memory: the cell's memory state
        [batchSize, memDim]

        control: the cell's control state
        [batchSize, ctrlDim]

    Returns the information extracted.
    [batchSize, memDim]
    '''
    def read(self, knowledgeBase, memory, control, name = "", reuse = None):
        with tf.variable_scope("read" + name, reuse = reuse):
            dim = config.memDim 

            ## memory dropout
            if config.memoryVariationalDropout:
                memory = ops.applyVarDpMask(memory, self.memDpMask, self.dropouts["memory"])
            else:
                memory = tf.nn.dropout(memory, self.dropouts["memory"])

            ## Step 1: knowledge base / memory interactions 
            # parameters for knowledge base and memory projection 
            proj = None
            if config.readProjInputs:
                proj = {"dim": config.attDim, "shared": config.readProjShared, "dropout": self.dropouts["read"] }
                dim = config.attDim

            # parameters for concatenating knowledge base elements
            concat = {"x": config.readMemConcatKB, "proj": config.readMemConcatProj}

            # compute interactions between knowledge base and memory
            interactions, interDim = ops.mul(x = knowledgeBase, y = memory, dim = config.memDim, 
                proj = proj, concat = concat, interMod = config.readMemAttType, name = "memInter")

            projectedKB = proj.get("x") if proj else None

            # project memory interactions back to hidden dimension
            if config.readMemProj:
                interactions = ops.linear(interactions, interDim, dim, act = config.readMemAct, 
                    name = "memKbProj")
            else: 
                dim = interDim

            ## Step 2: compute interactions with control
            if config.readCtrl:
                # compute interactions with control
                if config.ctrlDim != dim:
                    control = ops.linear(control, ctrlDim, dim, name = "ctrlProj")

                interactions, interDim = ops.mul(interactions, control, dim, 
                    interMod = config.readCtrlAttType, concat = {"x": config.readCtrlConcatInter}, 
                    name = "ctrlInter")

                # optionally concatenate knowledge base elements
                if config.readCtrlConcatKB:
                    if config.readCtrlConcatProj:
                        addedInp, addedDim = projectedKB, config.attDim
                    else:
                        addedInp, addedDim = knowledgeBase, config.memDim
                    interactions = tf.concat([interactions, addedInp], axis = -1)
                    dim += addedDim   

                # optional nonlinearity 
                interactions = ops.activations[config.readCtrlAct](interactions)

            ## Step 3: sum attentions up over the knowledge base
            # transform vectors to attention distribution
            attention = ops.inter2att(interactions, dim, dropout = self.dropouts["read"])

            self.attentions["kb"].append(attention)

            # optionally use projected knowledge base instead of original
            if config.readSmryKBProj:
                knowledgeBase = projectedKB
            
            # sum up the knowledge base according to the distribution
            information = ops.att2Smry(attention, knowledgeBase)

        return information 

    '''
    The write unit integrates newly retrieved information (from the read unit),
    with the cell's previous memory hidden state, resulting in a new memory value.
    The unit optionally supports:
    1. Self-attention to previous control / memory states, in order to consider previous steps
    in the reasoning process.
    2. Gating between the new memory and previous memory states, to allow dynamic adjustment
    of the reasoning process length.

    Args:
        memory: the cell's memory state
        [batchSize, memDim]

        info: the information to integrate with the memory
        [batchSize, memDim]

        control: the cell's control state
        [batchSize, ctrlDim]

        contControl: optional corresponding continuous control state 
        (before casting the attention over the words).
        [batchSize, ctrlDim]

    Return the new memory 
    [batchSize, memDim]
    '''
    def write(self, memory, info, control, contControl = None, name = "", reuse = None):
        with tf.variable_scope("write" + name, reuse = reuse):

            # optionally project info
            if config.writeInfoProj:
                info = ops.linear(info, config.memDim, config.memDim, name = "info")

            # optional info nonlinearity
            info = ops.activations[config.writeInfoAct](info) 

            # compute self-attention vector based on previous controls and memories
            if config.writeSelfAtt:
                selfControl = control
                if config.writeSelfAttMod == "CONT":
                    selfControl = contControl
                # elif config.writeSelfAttMod == "POST":
                #     selfControl = postControl
                selfControl = ops.linear(selfControl, config.ctrlDim, config.ctrlDim, name = "ctrlProj")
                
                interactions = self.controls * tf.expand_dims(selfControl, axis = 1)

                # if config.selfAttShareInter: 
                #     selfAttlogits = self.linearP(selfAttInter, config.encDim, 1, self.interL[0], self.interL[1], name = "modSelfAttInter")
                attention = ops.inter2att(interactions, config.ctrlDim, name = "selfAttention")
                self.attentions["self"].append(attention) 
                selfSmry = ops.att2Smry(attention, self.memories)

            # get write unit inputs: previous memory, the new info, optionally self-attention / control
            newMemory, dim = memory, config.memDim
            if config.writeInputs == "INFO":
                newMemory = info
            elif config.writeInputs == "SUM":
                newMemory += info
            elif config.writeInputs == "BOTH":
                newMemory, dim = ops.concat(newMemory, info, dim, mul = config.writeConcatMul)
            # else: MEM

            if config.writeSelfAtt:
                newMemory = tf.concat([newMemory, selfSmry], axis = -1)
                dim += config.memDim

            if config.writeMergeCtrl:
                newMemory = tf.concat([newMemory, control], axis = -1)
                dim += config.memDim                

            # project memory back to memory dimension
            if config.writeMemProj or (dim != config.memDim):
                newMemory = ops.linear(newMemory, dim, config.memDim, name = "newMemory")

            # optional memory nonlinearity
            newMemory = ops.activations[config.writeMemAct](newMemory)

            # write unit gate
            if config.writeGate:
                gateDim = config.memDim
                if config.writeGateShared:
                    gateDim = 1
                
                z = tf.sigmoid(ops.linear(control, config.ctrlDim, gateDim, name = "gate", bias = config.writeGateBias))
                
                self.attentions["gate"].append(z)

                newMemory = newMemory * z + memory * (1 - z)                

            # optional batch normalization
            if config.memoryBN:
                newMemory = tf.contrib.layers.batch_norm(newMemory, decay = config.bnDecay, 
                    center = config.bnCenter, scale = config.bnScale, 
                    is_training = self.train, updates_collections = None)

        return newMemory

    def memAutoEnc(newMemory, info, control, name = "", reuse = None):
        with tf.variable_scope("memAutoEnc" + name, reuse = reuse):
            # inputs to auto encoder
            features = info if config.autoEncMemInputs == "INFO" else newMemory
            features = ops.linear(features, config.memDim, config.ctrlDim, 
                act = config.autoEncMemAct, name = "aeMem")

            # reconstruct control
            if config.autoEncMemLoss == "CONT":
                loss = tf.reduce_mean(tf.squared_difference(control, features))
            else:                    
                interactions, dim = ops.mul(self.questionCntxWords, features, config.ctrlDim, 
                    concat = {"x": config.autoEncMemCnct}, mulBias = config.mulBias, name = "aeMem")
                
                logits = ops.inter2logits(interactions, dim)
                logits = self.expMask(logits, self.questionLengths)

                # reconstruct word attentions
                if config.autoEncMemLoss == "PROB":
                    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
                        labels = self.attentions["question"][-1], logits = logits))
                
                # reconstruct control through words attentions
                else:
                    attention = tf.nn.softmax(logits)
                    summary = ops.att2Smry(attention, self.questionCntxWords)
                    loss = tf.reduce_mean(tf.squared_difference(control, summary))
        
        return loss

    '''
    Call the cell to get new control and memory states.

    Args:
        inputs: in the current implementation the cell don't get recurrent inputs
        every iteration (argument for comparability with rnn interface).
            
        state: the cell current state (control, memory)
        MACCellTuple([batchSize, ctrlDim],[batchSize, memDim])

    Returns the new state -- the new memory and control values.
    MACCellTuple([batchSize, ctrlDim],[batchSize, memDim])
    '''
    def __call__(self, inputs, state, scope = None):
        scope = scope or type(self).__name__
        with tf.variable_scope(scope, reuse = self.reuse): #  as tfscope
            control = state.control
            memory = state.memory

            # cell sharing
            inputName = "qInput"
            inputNameU = "qInputU"
            inputReuseU = inputReuse = (self.iteration > 0)
            if config.controlInputUnshared:
                inputNameU = "qInput%d" % self.iteration
                inputReuseU = None

            cellName = ""
            cellReuse = (self.iteration > 0)
            if config.unsharedCells:
                cellName = str(self.iteration)
                cellReuse = None 

            ## control unit
            # prepare question input to control 
            controlInput = ops.linear(self.vecQuestions, config.ctrlDim, config.ctrlDim, 
                name = inputName, reuse = inputReuse)

            controlInput = ops.activations[config.controlInputAct](controlInput)

            controlInput = ops.linear(controlInput, config.ctrlDim, config.ctrlDim, 
                name = inputNameU, reuse = inputReuseU)

            newControl, self.contControl = self.control(controlInput, self.inWords, self.outWords, 
                self.questionLengths, control, self.contControl, name = cellName, reuse = cellReuse)
            
            # read unit
            # ablation: use whole question as control
            if config.controlWholeQ:                    
                newControl = self.vecQuestions
                # ops.linear(self.vecQuestions, config.ctrlDim, projDim, name = "qMod") 

            info = self.read(self.knowledgeBase, memory, newControl, name = cellName, reuse = cellReuse) 

            if config.writeDropout < 1.0:
                # write unit
                info = tf.nn.dropout(info, self.dropouts["write"])
            
            newMemory = self.write(memory, info, newControl, self.contControl, name = cellName, reuse = cellReuse)

            # add auto encoder loss for memory
            # if config.autoEncMem:
            #     self.autoEncLosses["memory"] += memAutoEnc(newMemory, info, newControl)
        
            # append as standard list?
            self.controls = tf.concat([self.controls, tf.expand_dims(newControl, axis = 1)], axis = 1)
            self.memories = tf.concat([self.memories, tf.expand_dims(newMemory, axis = 1)], axis = 1)
            self.infos = tf.concat([self.infos, tf.expand_dims(info, axis = 1)], axis = 1)

            # self.contControls = tf.concat([self.contControls, tf.expand_dims(contControl, axis = 1)], axis = 1)
            # self.postControls = tf.concat([self.controls, tf.expand_dims(postControls, axis = 1)], axis = 1)

        newState = MACCellTuple(newControl, newMemory)
        return self.none, newState

    '''
    Initializes the a hidden state to based on the value of the initType:
    "PRM" for parametric initialization
    "ZERO" for zero initialization  
    "Q" to initialize to question vectors.

    Args:
        name: the state variable name.
        dim: the dimension of the state.
        initType: the type of the initialization
        batchSize: the batch size

    Returns the initialized hidden state.
    '''
    def initState(self, name, dim, initType, batchSize):
        if initType == "PRM":
            prm = tf.get_variable(name, shape = (dim, ),
                    initializer = tf.random_normal_initializer())                
            initState = tf.tile(tf.expand_dims(prm, axis = 0), [batchSize, 1])
        elif initType == "ZERO":
            initState = tf.zeros((batchSize, dim), dtype = tf.float32)
        else: # "Q"
            initState = self.vecQuestions
        return initState

    '''
    Add a parametric null word to the questions.

    Args:
        words: the words to add a null word to.
        [batchSize, questionLentgth]

        lengths: question lengths.
        [batchSize] 

    Returns the updated word sequence and lengths.  
    '''
    def addNullWord(words, lengths):
        nullWord = tf.get_variable("zeroWord", shape = (1 , config.ctrlDim), initializer = tf.random_normal_initializer())                    
        nullWord = tf.tile(tf.expand_dims(nullWord, axis = 0), [self.batchSize, 1, 1])
        words = tf.concat([nullWord, words], axis = 1)
        lengths += 1
        return words, lengths

    '''
    Initializes the cell internal state (currently it's stateful). In particular,
    1. Data-structures (lists of attention maps and accumulated losses).
    2. The memory and control states.
    3. The knowledge base (optionally merging it with the question vectors)
    4. The question words used by the cell (either the original word embeddings, or the 
       encoder outputs, with optional projection).

    Args:
        batchSize: the batch size

    Returns the initial cell state.
    '''
    def zero_state(self, batchSize, dtype = tf.float32):
        ## initialize data-structures
        self.attentions = {"kb": [], "question": [], "self": [], "gate": []}
        self.autoEncLosses = {"control": tf.constant(0.0), "memory": tf.constant(0.0)}


        ## initialize state
        initialControl = self.initState("initCtrl", config.ctrlDim, config.initCtrl, batchSize)
        initialMemory = self.initState("initMem", config.memDim, config.initMem, batchSize)

        self.controls = tf.expand_dims(initialControl, axis = 1)
        self.memories = tf.expand_dims(initialMemory, axis = 1)
        self.infos = tf.expand_dims(initialMemory, axis = 1)
        
        self.contControl = initialControl
        # self.contControls = tf.expand_dims(initialControl, axis = 1)
        # self.postControls = tf.expand_dims(initialControl, axis = 1)


        ## initialize knowledge base
        # optionally merge question into knowledge base representation
        if config.initKBwithQ != "NON":
            iVecQuestions = ops.linear(self.vecQuestions, config.ctrlDim, config.memDim, name = "questions") 

            concatMul = (config.initKBwithQ == "MUL") 
            cnct, dim = ops.concat(self.knowledgeBase, iVecQuestions, config.memDim, mul = concatMul, expandY = True)
            self.knowledgeBase = ops.linear(cnct, dim, config.memDim, name = "initKB")


        ## initialize question words
        # choose question words to work with (original embeddings or encoder outputs)
        words = self.questionCntxWords if config.controlContextual else self.questionWords
        
        # optionally add parametric "null" word in the to all questions
        if config.addNullWord:
            words, questionLengths = self.addNullWord(words, questionLengths)

        # project words
        self.inWords = self.outWords = words
        if config.controlInWordsProj or config.controlOutWordsProj:
            pWords = ops.linear(words, config.ctrlDim, config.ctrlDim, name = "wordsProj") 
            self.inWords = pWords if config.controlInWordsProj else words
            self.outWords = pWords if config.controlOutWordsProj else words

        # if config.controlCoverage:
        #     self.coverage = tf.zeros((batchSize, tf.shape(words)[1]), dtype = tf.float32)
        #     self.coverageBias = tf.get_variable("coverageBias", shape = (),
        #         initializer = config.controlCoverageBias)  

        ## initialize memory variational dropout mask
        if config.memoryVariationalDropout:
            self.memDpMask = ops.generateVarDpMask((batchSize, config.memDim), self.dropouts["memory"])

        return MACCellTuple(initialControl, initialMemory)