Block Attention Residuals reduces depth memory by summarizing fixed groups of completed layers. The embedding is always source zero. Each completed block contributes the sum of its layer outputs, and a nonempty current block contributes the sum of outputs already produced inside that block.
For block size m, completed block j contains the chronological layer outputs from jm through (j+1)m−1:
Bj=i=jm∑(j+1)m−1hi.When the number of preceding layers is not divisible by m, sum the remaining outputs into one partial-block source. Do not add an empty partial source at an exact block boundary. A supplied helper performs the shared Full Attention Residuals read over the constructed sources. Return a dictionary with retrieved_representation, attention_weights, and block_sources. The tensors preserve the input floating-point dtype and device.
Examples
Input: embedding = [[[1,0]]], previous_outputs = [], pseudo_query = [1,0.5], block_size = 2, eps = 1e-06
Output: {"retrieved_representation": tensor([[[1.0, 0.0]]]), "attention_weights": tensor([[[1.0]]]), "block_sources": tensor([[[[1.0, 0.0]]]])}
Explanation: With no previous layer outputs, the embedding is the only block source, so it is returned unchanged with weight 1.
Input: embedding shape (1, 1, 2), previous outputs shape (3, 1, 1, 2), pseudo-query shape (2,), block size = 2
Output: {"retrieved_representation": tensor of shape (1, 1, 2), "attention_weights": tensor of shape (3, 1, 1), "block_sources": tensor of shape (3, 1, 1, 2)}
Input: embedding shape (1, 2, 2), previous outputs shape (4, 1, 2, 2), pseudo-query shape (2,), block size = 2
Output: {"retrieved_representation": tensor of shape (1, 2, 2), "attention_weights": tensor of shape (3, 1, 2), "block_sources": tensor of shape (3, 1, 2, 2)}
Hint 1
Use the layer-count remainder to separate complete blocks from a possible partial block.
Hint 2
Sum each complete slice along the layer dimension before stacking the sources.
Hint 3
Append the partial sum only when the remainder is nonzero, then call the supplied helper.
Requirements
- Return a dictionary with exactly retrieved_representation, attention_weights, and block_sources.
- Keep the embedding as the first source.
- Sum each consecutive group of block-size layer outputs into one completed source.
- Add one partial source only when ungrouped layer outputs remain.
- Preserve tensor dtype and device without mutating inputs.
Constraints
- Preceding outputs are ordered chronologically and may contain zero layers.
- Block size is a positive integer.
- Embedding and layer outputs share shape (batch, sequence, width).
- The pseudo-query has length equal to the feature width.
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Accepts: number
Block AttnRes
Block Attention Residuals reduces depth memory by summarizing fixed groups of completed layers. The embedding is always source zero. Each completed block contributes the sum of its layer outputs, and a nonempty current block contributes the sum of outputs already produced inside that block.
For block size m, completed block j contains the chronological layer outputs from jm through (j+1)m−1:
Bj=i=jm∑(j+1)m−1hi.When the number of preceding layers is not divisible by m, sum the remaining outputs into one partial-block source. Do not add an empty partial source at an exact block boundary. A supplied helper performs the shared Full Attention Residuals read over the constructed sources. Return a dictionary with retrieved_representation, attention_weights, and block_sources. The tensors preserve the input floating-point dtype and device.
Examples
Input: embedding = [[[1,0]]], previous_outputs = [], pseudo_query = [1,0.5], block_size = 2, eps = 1e-06
Output: {"retrieved_representation": tensor([[[1.0, 0.0]]]), "attention_weights": tensor([[[1.0]]]), "block_sources": tensor([[[[1.0, 0.0]]]])}
Explanation: With no previous layer outputs, the embedding is the only block source, so it is returned unchanged with weight 1.
Input: embedding shape (1, 1, 2), previous outputs shape (3, 1, 1, 2), pseudo-query shape (2,), block size = 2
Output: {"retrieved_representation": tensor of shape (1, 1, 2), "attention_weights": tensor of shape (3, 1, 1), "block_sources": tensor of shape (3, 1, 1, 2)}
Input: embedding shape (1, 2, 2), previous outputs shape (4, 1, 2, 2), pseudo-query shape (2,), block size = 2
Output: {"retrieved_representation": tensor of shape (1, 2, 2), "attention_weights": tensor of shape (3, 1, 2), "block_sources": tensor of shape (3, 1, 2, 2)}
Hint 1
Use the layer-count remainder to separate complete blocks from a possible partial block.
Hint 2
Sum each complete slice along the layer dimension before stacking the sources.
Hint 3
Append the partial sum only when the remainder is nonzero, then call the supplied helper.
Requirements
- Return a dictionary with exactly retrieved_representation, attention_weights, and block_sources.
- Keep the embedding as the first source.
- Sum each consecutive group of block-size layer outputs into one completed source.
- Add one partial source only when ungrouped layer outputs remain.
- Preserve tensor dtype and device without mutating inputs.
Constraints
- Preceding outputs are ordered chronologically and may contain zero layers.
- Block size is a positive integer.
- Embedding and layer outputs share shape (batch, sequence, width).
- The pseudo-query has length equal to the feature width.
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Accepts: number