Gated Cross-Attention (Flamingo tanh Gate)
Problem Statement
Flamingo inserts new cross-attention layers into a frozen LM and initializes them to be no-ops via a tanh gate that starts at zero. Implement the gated residual so the model behaves identically to the base LM at initialization.
Background
A gated cross-attention block adds its attention output back through a learnable scalar gate alpha:
y=x+tanh(α)⋅attn_out
At alpha = 0, tanh(0) = 0, so y = x exactly — the injected layer is invisible and the frozen LM is preserved. As training moves alpha away from 0 the visual pathway gradually opens. Given x, the precomputed attn_out, and alpha, return y.
Your Task
Implement:
def gated_cross_attention(x, attn_out, alpha):
Return y as a list rounded to 4 decimals.
Input Format
- x, attn_out: lists of the same length D.
- alpha (float): the gate pre-activation.
Output Format
- A list of D floats rounded to 4 decimals.
Sample
print(gated_cross_attention([1.0, 2.0], [10.0, 10.0], 0.0))
Output:
[1.0, 2.0]
Example:
print(gated_cross_attention([1.0, 2.0], [10.0, 10.0], 0.0))
[1.0, 2.0]
- Calculate the gate activation by applying the hyperbolic tangent to the input α: tanh(0.0)=0.0. This step determines the scaling factor for the attention output, ensuring the layer acts as a no-op at initialization.
- Scale the attention output vector by the gate value: [10.0,10.0]⋅0.0=[0.0,0.0]. This zeroes out the contribution of the cross-attention mechanism.
- Add the scaled attention output to the input vector x to compute the residual connection: [1.0,2.0]+[0.0,0.0]=[1.0,2.0]. This preserves the original input values since the gate is closed.
- Round each element of the resulting vector to 4 decimal places: [1.0,2.0] remains [1.0,2.0]. This ensures the output meets the required precision format.
- The final output is [1.0, 2.0]
Constraints:
1 <= D <= 4096.- The gate is
tanh(alpha); atalpha == 0the block is an exact identity. - Round every entry to 4 decimals; avoid
-0.0.
1. Background Knowledge
In Vision-Language Models (VLMs) like Flamingo, new cross-attention layers are inserted into a frozen Large Language Model (LLM). To prevent these new layers from disrupting the pre-trained LLM's behavior at the start of training, they are initialized as no-ops (identity operations). This is achieved using a gated residual connection where the attention output is scaled by a learnable scalar gate.
The core mechanism relies on the tanh activation function. The gate is defined as tanh(α), where α is a learnable parameter. At initialization, α is set to 0. Since tanh(0)=0, the scaled attention output becomes zero, and the residual connection y=x+0 ensures the output y is exactly equal to the input x. As training progresses, α moves away from zero, gradually "opening" the gate and allowing visual information to flow into the language model.
This pattern is a specific instance of gated linear units or residual gating, commonly used in architectures like Gated Attention and Highway Networks. The key insight is that by initializing the gate to zero, you can safely inject new parameters into a pre-trained model without catastrophic forgetting. The tanh function is chosen because it is smooth, bounded in [−1,1], and has a non-zero derivative at zero, allowing for stable gradient flow during the initial training steps.
2. Algorithm Approach
The problem requires implementing a simple element-wise operation combined with a scalar multiplication. The approach is:
- Compute the gate value: g=tanh(α).
- Scale the attention output: s=g⋅attn_out (element-wise).
- Add the scaled output to the input: y=x+s (element-wise).
- Round each element of y to 4 decimal places.
This is a vectorized operation that can be performed in a single pass over the input lists. No complex loops or data structures are needed beyond basic list comprehension or array operations.
3. Step-by-Step Strategy
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.