Problems
Loading...
Step 3 of 11
Full AttnRes
Medium

Full Attention Residuals treats the token embedding as depth source zero and every preceding layer output as another source. For pseudo-query wlw_l, source key kik_i, and raw source value viv_i, the score kernel is

ϕ(wl,ki)=exp⁡(wlTRMSNorm⁡(ki)).\phi(w_l,k_i)=\exp\left(w_l^{\mathsf T}\operatorname{RMSNorm}(k_i)\right).

For LL available depth sources, the normalized weight and retrieved representation are

αi→l=ϕ(wl,ki)∑j=0L−1ϕ(wl,kj).\alpha_{i\rightarrow l}=\frac{\phi(w_l,k_i)}{\sum_{j=0}^{L-1}\phi(w_l,k_j)}. hl=∑i=0L−1αi→lvi.h_l=\sum_{i=0}^{L-1}\alpha_{i\rightarrow l}v_i.

Keys are normalized before scoring, while values remain the original source representations. Return a dictionary with retrieved_representation and attention_weights. Both tensors preserve the input floating-point dtype and device.

Loading visualization...

Examples

Input: embedding = [[[2]]], previous_outputs = [], pseudo_query = [1], eps = 1e-06

Output: {"retrieved_representation": tensor([[[2.0]]]), "attention_weights": tensor([[[1.0]]])}

Explanation: With no previous layer outputs, the embedding is the only depth source, so its attention weight is 1.

Input: embedding = [[[1,0],[0.5,1]]], previous_outputs = [[[[0,2],[1.5,-0.5]]]], pseudo_query = [1,-0.5], eps = 1e-06

Output: {"retrieved_representation": tensor of shape (1, 2, 2), "attention_weights": tensor of shape (2, 1, 2)}

Input: embedding = [[[1,-1]]], previous_outputs = [[[[2,0.5]]],[[[-0.5,1.5]]]], pseudo_query = [0.75,1.25], eps = 1e-05

Output: {"retrieved_representation": tensor of shape (1, 1, 2), "attention_weights": tensor of shape (3, 1, 1)}

Hint 1

Prepend the embedding by adding a depth dimension before concatenating previous outputs.

Hint 2

Compute the root mean square scale along the feature dimension while keeping that dimension.

Hint 3

Apply softmax along the depth dimension before weighting the unnormalized sources.

Requirements

  • Return a dictionary with exactly retrieved_representation and attention_weights. Their shapes are (batch, sequence, width) and (sources, batch, sequence).
  • Depth weights sum to one for every batch item and token position.
  • Preserve the input floating-point dtype and device.
  • Do not mutate the embedding or preceding outputs.

Constraints

  • The embedding has shape (batch, sequence, width).
  • Preceding outputs have shape (layers, batch, sequence, width) and may contain zero layers.
  • The pseudo-query has shape (width,).
  • The normalization constant is positive.
Prev3 / 11Next