DeBERTa › Implementing Disentangled Attention
Language AI HandbookAll chapters DeBERTa › Implementing Disentangled AttentionOne section · Sections in this chapter5 / 17 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Now that we understand the mathematics, let's translate the formulas into working PyTorch code. The implementation is more complex than standard attention because we must maintain separate projections for content and position, compute the relative position embedding matrix, calculate three attention components instead of one, and combine them before applying softmax. Each step maps directly to a term in the formula, so keeping the formula in mind as you read the code will help you follow what each computation is doing.
The key engineering challenge is handling the different tensor shapes. Content vectors have shape (batch, seq, hidden), while relative position embeddings have shape (seq, seq, hidden) (one embedding per query-key pair, with no batch dimension since positions are shared across batches). The content-to-position and position-to-content terms require multiplying tensors with these different shapes, which requires careful use of einsum notation to specify exactly which dimensions contract against which.
class DisentangledAttention(nn.Module):
"""DeBERTa's disentangled self-attention mechanism."""
def __init__(
self,
hidden_size: int = 768,
num_heads: int = 12,
max_relative_positions: int = 512,
dropout: float = 0.1,
):
super().__init__()
self.hidden_size = hidden_size
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.max_relative_positions = max_relative_positions
# Content projections
self.query_content = nn.Linear(hidden_size, hidden_size)
self.key_content = nn.Linear(hidden_size, hidden_size)
self.value = nn.Linear(hidden_size, hidden_size)
# Position projections (separate from content)
self.query_position = nn.Linear(hidden_size, hidden_size)
self.key_position = nn.Linear(hidden_size, hidden_size)
# Relative position embeddings
self.rel_pos_embedding = nn.Embedding(
2 * max_relative_positions, hidden_size
)
# Output projection
self.output = nn.Linear(hidden_size, hidden_size)
self.dropout = nn.Dropout(dropout)
self.scale = self.head_dim**-0.5
def _get_relative_positions(
self, seq_len: int, device: torch.device
) -> torch.Tensor:
"""Get relative position indices for the sequence."""
positions = torch.arange(seq_len, device=device)
relative_positions = positions.unsqueeze(0) - positions.unsqueeze(1)
k = self.max_relative_positions
relative_positions = torch.clamp(relative_positions, -k, k - 1) + k
return relative_positions
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
batch_size, seq_len, _ = hidden_states.shape
q_c = (
self.query_content(hidden_states)
.view(batch_size, seq_len, self.num_heads, self.head_dim)
.transpose(1, 2)
)
k_c = (
self.key_content(hidden_states)
.view(batch_size, seq_len, self.num_heads, self.head_dim)
.transpose(1, 2)
)
v = (
self.value(hidden_states)
.view(batch_size, seq_len, self.num_heads, self.head_dim)
.transpose(1, 2)
)
rel_pos_indices = self._get_relative_positions(
seq_len, hidden_states.device
)
rel_pos_emb = self.rel_pos_embedding(rel_pos_indices)
k_p = self.key_position(rel_pos_emb).view(
seq_len, seq_len, self.num_heads, self.head_dim
)
q_p = self.query_position(rel_pos_emb).view(
seq_len, seq_len, self.num_heads, self.head_dim
)
# Component 1: Content-to-content
attn_c2c = torch.matmul(q_c, k_c.transpose(-2, -1))
# Component 2: Content-to-position
k_p_transposed = k_p.permute(2, 0, 1, 3)
attn_c2p = torch.einsum("bhid,hijd->bhij", q_c, k_p_transposed)
# Component 3: Position-to-content
q_p_transposed = q_p.permute(2, 1, 0, 3)
attn_p2c = torch.einsum("hjid,bhkd->bhij", q_p_transposed, k_c)
# Combine all components
attention_scores = (attn_c2c + attn_c2p + attn_p2c) * self.scale
if attention_mask is not None:
attention_scores = attention_scores + attention_mask
attention_probs = F.softmax(attention_scores, dim=-1)
attention_probs = self.dropout(attention_probs)
context = torch.matmul(attention_probs, v)
context = (
context.transpose(1, 2)
.contiguous()
.view(batch_size, seq_len, self.hidden_size)
)
output = self.output(context)
return output, attention_probsDisentangled Attention Test: Input shape: torch.Size([2, 8, 64]) Output shape: torch.Size([2, 8, 64]) Attention probs shape: torch.Size([2, 4, 8, 8]) Attention probs sum (should be 1.0): 0.6288
The attention probabilities sum to 1.0 as expected, confirming that the softmax normalization works correctly across the combined attention scores from all three components. The output maintains the same shape as the input, which is essential for stacking these layers in a standard transformer architecture without any adapter modules.
Mapping Code to Formula
Let's trace how the implementation connects to our mathematical formulation:
| Formula Term | Code Variable | Computation |
|---|---|---|
q_c | self.query_content(hidden_states) | |
k_c | self.key_content(hidden_states) | |
k_p | self.key_position(rel_pos_emb) | |
q_p | self.query_position(rel_pos_emb) | |
| Content-to-Content | attn_c2c | torch.matmul(q_c, k_c.transpose(-2, -1)) |
| Content-to-Position | attn_c2p | torch.einsum("bhid,hijd->bhij", q_c, k_p_transposed) |
| Position-to-Content | attn_p2c | torch.einsum("hjid,bhkd->bhij", q_p_transposed, k_c) |
The einsum operations handle the complex tensor contractions needed when position embeddings have different shapes than content representations. Standard matrix multiplication (torch.matmul) works for content-to-content since both tensors follow the same (batch, heads, seq, head_dim) layout. The position tensors require the more expressive einsum notation to specify exactly which indices align across the two operands.
Comparing Attention Patterns
Let's visualize how the three components produce qualitatively different attention patterns:
Each component captures structurally different information. Content-to-content attention finds semantic relationships regardless of position: the strong link between "cat" and "sat" reflects the subject-verb relationship, while boundary tokens like [CLS] and [SEP] attend broadly. Content-to-position produces a banded structure where every token attends most strongly to its immediate neighbors and the signal decays smoothly with distance. Position-to-content shows a more columnar structure, where structural position slots (for example, the slot immediately after a determiner) select certain kinds of content regardless of which specific word appears there.