Implementation:FlagOpen FlagEmbedding RetroMAE EnhancedDecoder
| Knowledge Sources | |
|---|---|
| Domains | Machine_Learning, Deep_Learning, Transformer_Architecture |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
Enhanced decoder architecture for RetroMAE with custom attention mechanisms supporting matrix attention masks.
Description
This module provides modified BERT-style attention layers that support RetroMAE's unique matrix attention masking strategy:
BertSelfAttention implements multi-head self-attention with separate query, key, and value inputs (instead of the standard single hidden state input). This enables the decoder to use different representations for attention queries versus keys/values, which is crucial for RetroMAE's masked auto-encoding objective.
BertAttention wraps BertSelfAttention and adds the output projection layer, maintaining compatibility with the standard BERT architecture while supporting the enhanced attention mechanism.
BertLayerForDecoder combines the enhanced attention with standard feed-forward layers and supports cross-attention for encoder-decoder architectures. The key modification is accepting separate query, key, and value tensors in the forward pass, allowing each decoder position to attend to customized masked versions of the input sequence via matrix attention masks.
These components enable RetroMAE's decoder to learn robust bidirectional representations by forcing each token to predict the full sequence from varied masked contexts.
Usage
Use these layers when building RetroMAE decoder architectures that require matrix attention masks for enhanced masked auto-encoding pre-training.
Code Reference
Source Location
- Repository: FlagOpen_FlagEmbedding
- File: research/baai_general_embedding/retromae_pretrain/enhancedDecoder.py
- Lines: 1-288
Signature
class BertSelfAttention(nn.Module):
def __init__(self, config, position_embedding_type=None)
def forward(self, query, key, value, attention_mask, ...)
class BertAttention(nn.Module):
def __init__(self, config, position_embedding_type=None)
def forward(self, query, key, value, attention_mask, ...)
class BertLayerForDecoder(nn.Module):
def __init__(self, config)
def forward(self, query, key, value, attention_mask, ...)
Import
from research.baai_general_embedding.retromae_pretrain.enhancedDecoder import BertLayerForDecoder
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| query | torch.Tensor | Yes | Query representations [batch, seq_len, hidden_dim] |
| key | torch.Tensor | Yes | Key representations [batch, seq_len, hidden_dim] |
| value | torch.Tensor | Yes | Value representations [batch, seq_len, hidden_dim] |
| attention_mask | torch.FloatTensor | No | Matrix attention mask [batch, seq_len, seq_len] |
| head_mask | torch.FloatTensor | No | Mask for attention heads |
| encoder_hidden_states | torch.FloatTensor | No | Hidden states for cross-attention |
| output_attentions | bool | No | Whether to return attention weights |
Outputs
| Name | Type | Description |
|---|---|---|
| hidden_states | torch.Tensor | Output hidden states [batch, seq_len, hidden_dim] |
| attention_probs | torch.Tensor | Attention weights (if output_attentions=True) |
| past_key_value | Tuple | Cached key/value for generation (if use_cache=True) |
Usage Examples
import torch
from transformers import BertConfig
from research.baai_general_embedding.retromae_pretrain.enhancedDecoder import BertLayerForDecoder
# Initialize decoder layer
config = BertConfig.from_pretrained("bert-base-uncased")
config.is_decoder = True
config.add_cross_attention = False
decoder_layer = BertLayerForDecoder(config)
# Prepare inputs with matrix attention mask
batch_size, seq_len, hidden_dim = 4, 128, 768
query = torch.randn(batch_size, seq_len, hidden_dim)
key = torch.randn(batch_size, seq_len, hidden_dim)
value = torch.randn(batch_size, seq_len, hidden_dim)
# Matrix attention mask: each row has a different mask pattern
# Shape: [batch, seq_len, seq_len]
matrix_attention_mask = torch.randint(0, 2, (batch_size, seq_len, seq_len)).float()
matrix_attention_mask = (1.0 - matrix_attention_mask) * -10000.0 # Convert to attention mask format
# Forward pass
outputs = decoder_layer(
query=query,
key=key,
value=value,
attention_mask=matrix_attention_mask
)
hidden_states = outputs[0] # [batch, seq_len, hidden_dim]