def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value, n_heads, n_attention_chunks, attention_type, dropout, mode): """Reversible transformer decoder layer. Args: d_model: int: depth of embedding d_ff: int: depth of feed-forward layer d_attention_key: int: depth of key vector for each attention head d_attention_value: int: depth of value vector for each attention head n_heads: int: number of attention heads n_attention_chunks: int: number of chunks for attention attention_type: class: attention class to use, such as DotProductAttention. dropout: float: dropout rate (how much to drop out) mode: str: 'train' or 'eval' Returns: the layer. """ pre_attention = [ Chunk(n_sections=n_attention_chunks), # pylint: disable=no-value-for-parameter tl.LayerNorm(), tl.Dup(), tl.Dup(), tl.Parallel( [ tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key) ], [ tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_key) ], [ tl.ComputeAttentionHeads(n_heads=n_heads, d_head=d_attention_value) ], ), ] attention = attention_type(mode=mode) # ReversibleAttentionHalfResidual requires that post_attention be linear in # its input (so the backward pass can be computed without knowing the input) post_attention = [ tl.ComputeAttentionOutput(n_heads=n_heads, d_model=d_model), Unchunk(n_sections=n_attention_chunks), # pylint: disable=no-value-for-parameter BroadcastedDropout(rate=dropout, mode=mode), # pylint: disable=no-value-for-parameter ] feed_forward = [ FeedForward(d_model, d_ff, dropout, mode=mode), ] return [ ReversibleAttentionHalfResidual(pre_attention, attention, post_attention), tl.ReversibleSwap(), ReversibleHalfResidual(feed_forward), tl.ReversibleSwap(), ]
def DecoderBlock(d_model, d_ff, d_attention_key, d_attention_value, n_heads, n_attention_chunks, attention_loop_stride, dropout, mode): """Reversible transformer decoder layer. Args: d_model: int: depth of embedding d_ff: int: depth of feed-forward layer d_attention_key: int: depth of key vector for each attention head d_attention_value: int: depth of value vector for each attention head n_heads: int: number of attention heads n_attention_chunks: int: number of chunks for attention attention_loop_stride: int: number of query elements to compute attention for in parallel. Set to 0 to disable memory-efficient attention. dropout: float: dropout rate (how much to drop out) mode: str: 'train' or 'eval' Returns: the layer. """ pre_attention = [ Chunk(n_sections=n_attention_chunks), # pylint: disable=no-value-for-parameter tl.LayerNorm(), tl.Dup(), tl.Dup(), tl.Parallel( [tl.Dense(d_attention_key * n_heads), SplitHeads(n_heads=n_heads)], # pylint: disable=no-value-for-parameter [tl.Dense(d_attention_key * n_heads), SplitHeads(n_heads=n_heads)], # pylint: disable=no-value-for-parameter [ tl.Dense(d_attention_value * n_heads), SplitHeads(n_heads=n_heads) ], # pylint: disable=no-value-for-parameter ), ] # TODO(kitaev): add dropout if attention_loop_stride < 1: # Use the standard implementation if no loop_stride is provided. attention = DotProductAttention(dropout=None, mode=mode) else: attention = MemoryEfficientDotProductAttention( loop_stride=attention_loop_stride, dropout=None, mode=mode) # ReversibleAttentionHalfResidual requires that post_attention be linear in # its input (so the backward pass can be computed without knowing the input) post_attention = [ JoinHeads(), # pylint: disable=no-value-for-parameter tl.Dense(d_model), Unchunk(n_sections=n_attention_chunks), # pylint: disable=no-value-for-parameter ] feed_forward = [ FeedForward(d_model, d_ff, dropout, mode=mode), ] return [ ReversibleAttentionHalfResidual(pre_attention, attention, post_attention), tl.ReversibleSwap(), ReversibleHalfResidual(feed_forward), tl.ReversibleSwap(), ]