Parameter Bottlenecks

scalar_mixing_weights.py

import tensorflow as tf

def mix_models_scalar(linear_outputs, rnn_outputs, lstm_outputs):
    # Generate scalar valued weights to combine outputs of a
    # Linear RNN, vanilla RNN and LSTM. Each of the inputs to
    # this function is a list of outputs from the respective model.
    
    context_window_size = 20
    all_model_outputs = [linear_outputs, rnn_outputs, lstm_outputs]
    
    # Get scalar valued weights for each model.
    mixing_weights = [tf.get_variable(model, [1.0],trainable=True)
               for model in ['linear_weight', 'rnn_weight', 'lstm_weight']]
    
    all_weighted_outputs = []
    # Multiply the relevant weight with every model's states
    for weight, model_output in zip(weights, all_model_outputs):
        weighted_output = [weight * model_output[i] for i in range(context_window_size)]
        all_weighted_outputs.append(weighted_output)

    # Add relevant weighted states from each model together.
    combined_outputs = [tf.add_n(*context_step) for context_step in zip(*all_weighted_outputs)]

    # This is a list of 20 (batch_size, hidden_dimension) tensors, 
    # having combined the weighted outputs of all three models.
    return combined_outputs