Skip to content

dual_stream_attention

dual_stream_attention

Dual-stream attention block: per-stream norm + joint attention.

DualStreamAttentionBlock 

DualStreamAttentionBlock(joint_attention, attention_normalization_primary, attention_normalization_secondary, dropout=0.1)

Bases: DualStreamBlock

Dual-stream joint attention block.

Both streams have independent normalization and share attention through joint K/V concatenation. Supports optional conditioning via adaptive normalization.

Initialize DualStreamAttentionBlock.

Parameters:

Name Type Description Default
joint_attention JointAttention

JointAttention module for dual-stream attention.

 required
attention_normalization_primary BlockNormalization

Normalization for primary stream.

 required
attention_normalization_secondary BlockNormalization

Normalization for secondary stream.

 required
dropout float

Dropout rate for attention residual connections.

 0.1
Source code in src/versatil/models/layers/transformer/block/dual_stream_attention.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def __init__(
    self,
    joint_attention: JointAttention,
    attention_normalization_primary: BlockNormalization,
    attention_normalization_secondary: BlockNormalization,
    dropout: float = 0.1,
):
    """Initialize DualStreamAttentionBlock.

    Args:
        joint_attention: JointAttention module for dual-stream attention.
        attention_normalization_primary: Normalization for primary stream.
        attention_normalization_secondary: Normalization for secondary stream.
        dropout: Dropout rate for attention residual connections.
    """
    super().__init__(
        attention_normalization_primary=attention_normalization_primary,
        dropout=dropout,
    )
    self.joint_attention = joint_attention
    self.attention_normalization_secondary = attention_normalization_secondary
    self.attention_dropout_secondary = nn.Dropout(dropout)

forward 

forward(hidden_states_primary, hidden_states_secondary, conditioning=None, attention_mask_primary=None, attention_mask_secondary=None, joint_attention_mask=None, positional_encoding_primary=None, positional_encoding_secondary=None)

Forward pass through dual-stream attention block.

Parameters:

Name Type Description Default
hidden_states_primary Tensor

Primary stream tokens (B, S, D_p).

 required
hidden_states_secondary Tensor

Secondary stream tokens (B, T, D_s).

 required
conditioning Tensor | None

Conditioning vector for adaptive normalization (B, C).

 None
attention_mask_primary Tensor | None

Padding mask (B, S), True = masked.

 None
attention_mask_secondary Tensor | None

Padding mask (B, T), True = masked.

 None
joint_attention_mask Tensor | None

Pre-built joint mask (B, 1, S+T, S+T).

 None
positional_encoding_primary RotaryPositionalEncoding | None

Optional RoPE for primary stream.

 None
positional_encoding_secondary RotaryPositionalEncoding | None

Optional RoPE for secondary stream.

 None

Returns:

Type Description
tuple[Tensor, Tensor]

Tuple of (primary_output (B, S, D_p), secondary_output (B, T, D_s)).
Source code in src/versatil/models/layers/transformer/block/dual_stream_attention.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def forward(
    self,
    hidden_states_primary: torch.Tensor,
    hidden_states_secondary: torch.Tensor,
    conditioning: torch.Tensor | None = None,
    attention_mask_primary: torch.Tensor | None = None,
    attention_mask_secondary: torch.Tensor | None = None,
    joint_attention_mask: torch.Tensor | None = None,
    positional_encoding_primary: RotaryPositionalEncoding | None = None,
    positional_encoding_secondary: RotaryPositionalEncoding | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Forward pass through dual-stream attention block.

    Args:
        hidden_states_primary: Primary stream tokens (B, S, D_p).
        hidden_states_secondary: Secondary stream tokens (B, T, D_s).
        conditioning: Conditioning vector for adaptive normalization (B, C).
        attention_mask_primary: Padding mask (B, S), True = masked.
        attention_mask_secondary: Padding mask (B, T), True = masked.
        joint_attention_mask: Pre-built joint mask (B, 1, S+T, S+T).
        positional_encoding_primary: Optional RoPE for primary stream.
        positional_encoding_secondary: Optional RoPE for secondary stream.

    Returns:
        Tuple of (primary_output (B, S, D_p), secondary_output (B, T, D_s)).
    """
    residual_primary = hidden_states_primary
    residual_secondary = hidden_states_secondary
    normed_primary, gate_primary = self.attention_normalization_primary(
        x=hidden_states_primary, condition=conditioning
    )
    normed_secondary, gate_secondary = self.attention_normalization_secondary(
        x=hidden_states_secondary, condition=conditioning
    )
    attention_output_primary, attention_output_secondary = self.joint_attention(
        hidden_states_primary=normed_primary,
        hidden_states_secondary=normed_secondary,
        attention_mask_primary=attention_mask_primary,
        attention_mask_secondary=attention_mask_secondary,
        joint_attention_mask=joint_attention_mask,
        positional_encoding_primary=positional_encoding_primary,
        positional_encoding_secondary=positional_encoding_secondary,
    )
    hidden_states_primary = self._apply_primary_attention_residual(
        residual=residual_primary,
        attention_output=attention_output_primary,
        gate=gate_primary,
    )
    hidden_states_secondary = (
        residual_secondary
        + gate_secondary
        * self.attention_dropout_secondary(attention_output_secondary)
    )
    return hidden_states_primary, hidden_states_secondary