class_docstr: A stack of layers for FunnelTransformer.
class_docstr_tok: ["A","stack","of","layers","for","FunnelTransformer","."]
docstr: Create a LayerStack specialized for FunnelTransformer.
  
  The design of this class follows the transformer.LayerStack. See the
  docstring of that class for how the layer stack is built. Here we only
  discuss the features unique to the Funnel Transformer.
  
  This implementation has subtle difference from the Funnel Transformer
  introduced in https://arxiv.org/abs/2006.03236.
  
  1. Application to encoder-decoder model.
  The original Funnel Transformer was proposed for the encoder-only
  architectures such as BERT. In A.3 section of the paper, they discuss
  potential extension of the core idea to other model architectures. For
  encoder-decoder models, the authors suggest that the Funnel Transformer idea
  can be used to modify the encoder such that "the key difference compared to
  conventional models is the source side compression Funnel-Transformer
  provides".
  
  Therefore, we don't modify the decoder, i.e., we use the standard
  transformer.LayerStack and this class is only applicable to the encoder.
  
  2. Relative attention
  We use the simplified reletive attention scalar from the T5, whereas the
  Funnel Transformer paper uses the relative attention from the Transformer-XL
  (https://arxiv.org/abs/1901.02860).
  
  3. The order of pooling operation
  In the Funnel Transformer paper, only the query is pooled while key and
  value are kept intact. The resulting attention output has the same length as
  the query, enabling the residual connection.
  
  In our implementation, we apply the regular SelfAttention and then apply the
  pooling to the output. Since each sequence position in query is
  independently computed, we expect the difference between these
  implmentations to be negligible.
  
  Args:
    layers: a list of TransformerLayer
    n_blocks: an integer specifying the number of Funnel Transformer blocks.
    block_param_size: a list of integers specifying the number of layers in
      each block.
    block_repeat_size: a list of integers specifying the number of repeated
      layers in each block. The repeated layers share the parameters.
    pooling_size: an integer specifying the pool size
    sublayers_initial: an optional list of sublayer functions
    sublayers_per_layer: an optional list of sublayer functions
    sublayers_final: an optional list of sublayer functions
    pooling_type: a string specifying the pooling type. One of "mean", "max",
      or "min".
    n_submodules: an integer specifying the number of submodules (e.g.,
      SelfAttention and DenseReluDense for each layer of a block.
docstr_tok: ["Create","a","LayerStack","specialized","for","FunnelTransformer",".","The","design","of","this","class","follows","the","transformer.LayerStack",".","See","the","docstring","of","that","class","for","how","the","layer","stack","is","built",".","Here","we","only","discuss","the","features","unique","to","the","Funnel","Transformer",".","This","implementation","has","subtle","difference","from","the","Funnel","Transformer","introduced","in","https",":","\/\/arxiv.org\/abs\/2006.03236",".","1",".","Application","to","encoder-decoder","model",".","The","original","Funnel","Transformer","was","proposed","for","the","encoder-only","architectures","such","as","BERT",".","In","A.3","section","of","the","paper",",","they","discuss","potential","extension","of","the","core","idea","to","other","model","architectures",".","For","encoder-decoder","models",",","the","authors","suggest","that","the","Funnel","Transformer","idea","can","be","used","to","modify","the","encoder","such","that","``","the","key","difference","compared","to","conventional","models","is","the","source","side","compression","Funnel-Transformer","provides","''",".","Therefore",",","we","do","n't","modify","the","decoder",",","i.e.",",","we","use","the","standard","transformer.LayerStack","and","this","class","is","only","applicable","to","the","encoder",".","2",".","Relative","attention","We","use","the","simplified","reletive","attention","scalar","from","the","T5",",","whereas","the","Funnel","Transformer","paper","uses","the","relative","attention","from","the","Transformer-XL","(","https",":","\/\/arxiv.org\/abs\/1901.02860",")",".","3",".","The","order","of","pooling","operation","In","the","Funnel","Transformer","paper",",","only","the","query","is","pooled","while","key","and","value","are","kept","intact",".","The","resulting","attention","output","has","the","same","length","as","the","query",",","enabling","the","residual","connection",".","In","our","implementation",",","we","apply","the","regular","SelfAttention","and","then","apply","the","pooling","to","the","output",".","Since","each","sequence","position","in","query","is","independently","computed",",","we","expect","the","difference","between","these","implmentations","to","be","negligible",".","Args",":","layers",":","a","list","of","TransformerLayer","n_blocks",":","an","integer","specifying","the","number","of","Funnel","Transformer","blocks",".","block_param_size",":","a","list","of","integers","specifying","the","number","of","layers","in","each","block",".","block_repeat_size",":","a","list","of","integers","specifying","the","number","of","repeated","layers","in","each","block",".","The","repeated","layers","share","the","parameters",".","pooling_size",":","an","integer","specifying","the","pool","size","sublayers_initial",":","an","optional","list","of","sublayer","functions","sublayers_per_layer",":","an","optional","list","of","sublayer","functions","sublayers_final",":","an","optional","list","of","sublayer","functions","pooling_type",":","a","string","specifying","the","pooling","type",".","One","of","``","mean","''",",","``","max","''",",","or","``","min","''",".","n_submodules",":","an","integer","specifying","the","number","of","submodules","(","e.g.",",","SelfAttention","and","DenseReluDense","for","each","layer","of","a","block","."]
code: def __init__(
      self,
      layers,
      n_blocks=gin.REQUIRED,
      block_param_size=gin.REQUIRED,
      block_repeat_size=gin.REQUIRED,
      pooling_size=2,
      sublayers_initial=None,
      sublayers_per_layer=None,
      sublayers_final=None,
      pooling_type="mean",
      n_submodules=2,
  ):
      """Create a LayerStack specialized for FunnelTransformer.
  
      The design of this class follows the transformer.LayerStack. See the
      docstring of that class for how the layer stack is built. Here we only
      discuss the features unique to the Funnel Transformer.
  
      This implementation has subtle difference from the Funnel Transformer
      introduced in https://arxiv.org/abs/2006.03236.
  
      1. Application to encoder-decoder model.
      The original Funnel Transformer was proposed for the encoder-only
      architectures such as BERT. In A.3 section of the paper, they discuss
      potential extension of the core idea to other model architectures. For
      encoder-decoder models, the authors suggest that the Funnel Transformer idea
      can be used to modify the encoder such that "the key difference compared to
      conventional models is the source side compression Funnel-Transformer
      provides".
  
      Therefore, we don't modify the decoder, i.e., we use the standard
      transformer.LayerStack and this class is only applicable to the encoder.
  
      2. Relative attention
      We use the simplified reletive attention scalar from the T5, whereas the
      Funnel Transformer paper uses the relative attention from the Transformer-XL
      (https://arxiv.org/abs/1901.02860).
  
      3. The order of pooling operation
      In the Funnel Transformer paper, only the query is pooled while key and
      value are kept intact. The resulting attention output has the same length as
      the query, enabling the residual connection.
  
      In our implementation, we apply the regular SelfAttention and then apply the
      pooling to the output. Since each sequence position in query is
      independently computed, we expect the difference between these
      implmentations to be negligible.
  
      Args:
        layers: a list of TransformerLayer
        n_blocks: an integer specifying the number of Funnel Transformer blocks.
        block_param_size: a list of integers specifying the number of layers in
          each block.
        block_repeat_size: a list of integers specifying the number of repeated
          layers in each block. The repeated layers share the parameters.
        pooling_size: an integer specifying the pool size
        sublayers_initial: an optional list of sublayer functions
        sublayers_per_layer: an optional list of sublayer functions
        sublayers_final: an optional list of sublayer functions
        pooling_type: a string specifying the pooling type. One of "mean", "max",
          or "min".
        n_submodules: an integer specifying the number of submodules (e.g.,
          SelfAttention and DenseReluDense for each layer of a block.
      """
      if len(block_param_size) != n_blocks:
          raise ValueError(
              "Number of blocks should match the length of block_param_size."
          )
  
      if len(block_repeat_size) != n_blocks:
          raise ValueError(
              "Number of blocks should match the length of block_repeat_size."
          )
  
      if len(layers) != sum(block_param_size) * n_submodules:
          raise ValueError(
              "Total number of submodules should match the number of layers."
          )
  
      self._layers = layers
      self.n_blocks = n_blocks
      self.block_param_size = block_param_size
      self.block_repeat_size = block_repeat_size
      self.pooling_size = pooling_size
  
      self._sublayers_initial = sublayers_initial
      self._sublayers_per_layer = sublayers_per_layer
      self._sublayers_final = sublayers_final
  
      if pooling_type == "mean":
          self.pool_fn = mtf.reduce_mean
      elif pooling_type == "max":
          self.pool_fn = mtf.reduce_max
      elif pooling_type == "min":
          self.pool_fn = mtf.reduce_min
      else:
          raise ValueError(
              "Unknown pooling type. Choose among 'mean', 'max' or 'min'"
          )
      self.n_submodules = n_submodules
  
file_path: mesh/mesh_tensorflow/transformer/funnel_transformer.py
imports: import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 1 to create a new FunnelTransformerLayerStack object from inherited base classes: transformer with example: obj = FunnelTransformerLayerStack(layers, n_blocks, block_param_size, block_repeat_size, pooling_size, sublayers_initial, sublayers_per_layer, sublayers_final, pooling_type, n_submodules)