beam_width & maximum_iterations in BeamSearchDecoder.call instead of init?

[StarWang]

Describe the feature and the current behavior/state.
Hi TFA team, the current seq2seq beam search decoder implementation takes in beam_width & maximum_iterations during decoder initialization init() instead of dynamic calls call. This makes it difficult for tuning beam_width when there's a trained model. However, it's a common request that in practice, when we have a trained model, we would like to tune the beam_width & maximum_iterations for best inference latency and performance without retraining or tweaking the model.

Can you support this request? Or is there any way to fulfill the above request in a simple way?

Relevant information

• Are you willing to contribute it (yes/no): no
• Are you willing to maintain it going forward? (yes/no): no
• Is there a relevant academic paper? (if so, where): no
• Is there already an implementation in another framework? (if so, where): no
• Was it part of tf.contrib? (if so, where): no

Which API type would this fall under (layer, metric, optimizer, etc.)
Addons - Seq2seq

Who will benefit with this feature?
TensorFlow Addons users who want to tune beam_width with input params beam_width during inference time (TF serving, JNI etc.), instead of tweaking models

Any other info.

[WindQAQ]

/cc @guillaumekln for visibility

[guillaumekln]

Did you try creating the BeamSearchDecoder instance in your call method?

I think having beam_width in the constructor is the correct design. But we could look to make it work with a dynamic value (i.e. a tensor).

[StarWang]

Hi @guillaumekln and @WindQAQ ,

Thanks for the issue edition and reply. I will try but maybe not in the short term future as currently I'm not focused on this.

I'm okay with keeping beam_width in the constructor, since you don't have to specify it in inference time if you are convinced about the value of beam_width. Just want to know if there's a way for dynamic tuning since it's common and very much needed to my understanding :)

[abhishek-niranjan]

Hi @StarWang

Hi TFA team, the current seq2seq beam search decoder implementation takes in beam_width & maximum_iterations during decoder initialization init() instead of dynamic calls call()
This makes it difficult for tuning beam_width when there's a trained model. However, it's a common request that in practice, when we have a trained model, we would like to tune the beam_width & maximum_iterations for best inference latency and performance without retraining or tweaking the model.

Though this is true, there's still a way to tune your beam_width in current implementation after your model has been trained which involves multiple instantiations of tfa.seq2seq.BeamSearchDecoder. For example,

Say the decoder stack looks like this:

class Decoder(tf.keras.Model):
  def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz, attention_type='luong'):
    super(Decoder, self).__init__()
    self.batch_sz = batch_sz
    self.dec_units = dec_units
    self.attention_type = attention_type
    
    # Embedding Layer
    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
    
    #Final Dense layer on which softmax will be applied
    self.fc = tf.keras.layers.Dense(vocab_size)

    # Define the fundamental cell for decoder recurrent structure
    self.decoder_rnn_cell = tf.keras.layers.LSTMCell(self.dec_units)
   
    # Sampler
    self.sampler = tfa.seq2seq.sampler.TrainingSampler()

    # Create attention mechanism with memory = None
    self.attention_mechanism = self.build_attention_mechanism(self.dec_units, 
                                                              None, self.batch_sz*[max_length_input], self.attention_type)

    # Wrap attention mechanism with the fundamental rnn cell of decoder
    self.rnn_cell = self.build_rnn_cell(batch_sz)

    # Define the decoder with respect to fundamental rnn cell
    self.decoder = tfa.seq2seq.BasicDecoder(self.rnn_cell, sampler=self.sampler, output_layer=self.fc)

    
  def build_rnn_cell(self, batch_sz):
    rnn_cell = tfa.seq2seq.AttentionWrapper(self.decoder_rnn_cell, 
                                  self.attention_mechanism, attention_layer_size=self.dec_units)
    return rnn_cell

  def build_attention_mechanism(self, dec_units, memory, memory_sequence_length, attention_type='luong'):
    if(attention_type=='bahdanau'):
      return tfa.seq2seq.BahdanauAttention(units=dec_units, memory=memory, memory_sequence_length=memory_sequence_length)
    else:
      return tfa.seq2seq.LuongAttention(units=dec_units, memory=memory, memory_sequence_length=memory_sequence_length)

  def build_initial_state(self, batch_sz, encoder_state, Dtype):
    decoder_initial_state = self.rnn_cell.get_initial_state(batch_size=batch_sz, dtype=Dtype)
    decoder_initial_state = decoder_initial_state.clone(cell_state=encoder_state)
    return decoder_initial_state


  def call(self, inputs, initial_state):
    x = self.embedding(inputs)
    outputs, _, _ = self.decoder(x, initial_state=initial_state, sequence_length=self.batch_sz*[max_length_output-1])
    return outputs

Then once you've trained your model, to test out outputs related to different beam widths, you can instantiate the BeamSearchDecoder multiple times without having to retrain your model.

# Define a beam-search decoder
decoder = tfa.seq2seq.BeamSearchDecoder(
    cell=decoderStack.rnn_cell,
    beam_width=beam_width,
    output_layer=projection_layer,
    length_penalty_weight=0.0,
    coverage_penalty_weight=0.0)

# decoding
outputs, _ = decoder(
    embedding_decoder,
    start_tokens=start_tokens,
    end_token=end_token,
    initial_state=decoder_initial_state)

Maybe it's not the ideal way to do it, but it serves the purpose of fine-tuning a beam_width and have future predictions according to it. Let me know if I'm missing something.

[guillaumekln]

Let's close this issue as @abhishek-niranjan shared a good way to that. @StarWang Feel free to reopen if you think something is missing to properly tune the beam width.

[StarWang]

Thanks all. I'll update this once I come across related issue.