From b3adf4bb5292a078de9ec735204e64e4b719658f Mon Sep 17 00:00:00 2001 From: Jingjing Xu Date: Mon, 13 Apr 2026 04:28:36 -0400 Subject: [PATCH 1/4] update --- i6_models/parts/conformer/mhsa_rel_pos.py | 28 ++++++++++++++--------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/i6_models/parts/conformer/mhsa_rel_pos.py b/i6_models/parts/conformer/mhsa_rel_pos.py index 6f05dc26..f7a906d5 100644 --- a/i6_models/parts/conformer/mhsa_rel_pos.py +++ b/i6_models/parts/conformer/mhsa_rel_pos.py @@ -163,15 +163,22 @@ def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> to k = key_seq.view(batch_dim_size, -1, self.num_heads, self.embed_dim_per_head) # [B, T', #heads, F'] if self.learnable_pos_emb: - pos_seq_q = torch.arange(time_dim_size, device=input_tensor.device) - pos_seq_k = torch.arange(time_dim_size, device=input_tensor.device) - - distance_mat = pos_seq_k[None, :] - pos_seq_q[:, None] - distance_mat_clipped = torch.clamp(distance_mat, -self.rel_pos_clip, self.rel_pos_clip) - - final_mat = distance_mat_clipped + self.rel_pos_clip - - rel_pos_embeddings = self.rel_pos_embeddings[final_mat] # [T, T', pos_emb_dim] + kv_pos_vec = torch.arange(time_dim_size, device=input_tensor.device) # [kv_len] + + query_spatial_dim_m1 = time_dim_size - 1 + q_pos_vec = torch.arange(query_spatial_dim_m1, device=input_tensor.device) # [q_len-1] + + # The min value is with kv_pos=0, q_pos=q_len-1: -(q_len-1) + # The max value is with kv_pos=kv_len-1, q_pos=0: k_len-1 + indices = torch.concat((q_pos_vec - query_spatial_dim_m1, kv_pos_vec), dim=-1) + indices = torch.clamp(indices, -self.rel_pos_clip, self.rel_pos_clip) + # Shift values to be >= 0. Each integer still uniquely identifies a relative position difference. + indices = indices + self.rel_pos_clip + rel_pos_embeddings = self.rel_pos_embeddings[indices] # [out_spatial_dim, n_out] + rel_pos_embeddings = rel_pos_embeddings.unsqueeze(0) + assert rel_pos_embeddings.shape == (1, 2 * time_dim_size - 1, self.pos_emb_dim), ( + "Something went wrong in reshaping" + ) else: rel_pos_embeddings = ( self._sinusoidal_pe( @@ -207,8 +214,7 @@ def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> to q_with_bias_v, rel_pos_embeddings.to(device=q_with_bias_v.device, dtype=q_with_bias_v.dtype), ) # [B, #heads, T, T'] or [B, #heads, T, T+T'+1] - if not self.learnable_pos_emb: - attn_bd = self._rel_shift_bhij(attn_bd, k_len=time_dim_size) # [B, #heads, T, T'] + attn_bd = self._rel_shift_bhij(attn_bd, k_len=time_dim_size) # [B, #heads, T, T'] # We use attn_mask to add BD matrix to attention scores. # From 74745ac2e03ef8d97fd1be15fa0ec505dcb3e4b1 Mon Sep 17 00:00:00 2001 From: Jingjing Xu Date: Fri, 17 Apr 2026 09:14:47 -0400 Subject: [PATCH 2/4] update --- i6_models/parts/conformer/mhsa_rel_pos.py | 82 ++++++++--------------- 1 file changed, 28 insertions(+), 54 deletions(-) diff --git a/i6_models/parts/conformer/mhsa_rel_pos.py b/i6_models/parts/conformer/mhsa_rel_pos.py index f7a906d5..820a5fdf 100644 --- a/i6_models/parts/conformer/mhsa_rel_pos.py +++ b/i6_models/parts/conformer/mhsa_rel_pos.py @@ -11,8 +11,8 @@ import torch.nn.functional as F from i6_models.config import ModelConfiguration -from i6_models.parts.dropout import BroadcastDropout from i6_models.util import compat +from i6_models.parts.dropout import BroadcastDropout @dataclass @@ -163,33 +163,19 @@ def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> to k = key_seq.view(batch_dim_size, -1, self.num_heads, self.embed_dim_per_head) # [B, T', #heads, F'] if self.learnable_pos_emb: - kv_pos_vec = torch.arange(time_dim_size, device=input_tensor.device) # [kv_len] - - query_spatial_dim_m1 = time_dim_size - 1 - q_pos_vec = torch.arange(query_spatial_dim_m1, device=input_tensor.device) # [q_len-1] - - # The min value is with kv_pos=0, q_pos=q_len-1: -(q_len-1) - # The max value is with kv_pos=kv_len-1, q_pos=0: k_len-1 - indices = torch.concat((q_pos_vec - query_spatial_dim_m1, kv_pos_vec), dim=-1) - indices = torch.clamp(indices, -self.rel_pos_clip, self.rel_pos_clip) - # Shift values to be >= 0. Each integer still uniquely identifies a relative position difference. - indices = indices + self.rel_pos_clip - rel_pos_embeddings = self.rel_pos_embeddings[indices] # [out_spatial_dim, n_out] - rel_pos_embeddings = rel_pos_embeddings.unsqueeze(0) - assert rel_pos_embeddings.shape == (1, 2 * time_dim_size - 1, self.pos_emb_dim), ( - "Something went wrong in reshaping" - ) - else: - rel_pos_embeddings = ( - self._sinusoidal_pe( - torch.arange( - time_dim_size - 1, -time_dim_size, -1, device=input_tensor.device, dtype=torch.float32 - ), - self.pos_emb_dim, - ) - .to(input_tensor.dtype) - .view(1, 2 * time_dim_size - 1, self.pos_emb_dim) + # 1D optimization: 2T-1 unique relative positions instead of T×T distance matrix. + # Descending order matches sinusoidal branch convention for _rel_shift_bhij. + rel_pos = torch.arange(time_dim_size - 1, -time_dim_size, -1, device=input_tensor.device) + # _rel_shift_bhij produces i-j convention; embedding table uses j-i, so negate + indices = torch.clamp(-rel_pos, -self.rel_pos_clip, self.rel_pos_clip) + self.rel_pos_clip + rel_pos_embeddings = self.rel_pos_embeddings[indices].view( + 1, 2 * time_dim_size - 1, self.pos_emb_dim ) # [1, T+T'-1, pos_emb_dim] + else: + rel_pos_embeddings = self._sinusoidal_pe( + torch.arange(time_dim_size - 1, -time_dim_size, -1, device=input_tensor.device, dtype=torch.float32), + self.pos_emb_dim, + ).view(1, 2 * time_dim_size - 1, self.pos_emb_dim) # [1, T+T'-1, pos_emb_dim] # dropout relative positional embeddings rel_pos_embeddings = self.pos_emb_dropout( @@ -208,42 +194,30 @@ def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> to q_with_bias_u = q + self.pos_bias_u if self.with_pos_bias else q # [B, T, #heads, F'] q_with_bias_v = q + self.pos_bias_v if self.with_pos_bias else q + # attention matrix a and c + attn_ac = torch.einsum("bihf, bjhf -> bhij", q_with_bias_u, k) # [B, #heads, T, T'] + # attention matrix b and d attn_bd = torch.einsum( - "bihf, ijhf -> bhij", - q_with_bias_v, - rel_pos_embeddings.to(device=q_with_bias_v.device, dtype=q_with_bias_v.dtype), + "bihf, ijhf -> bhij", q_with_bias_v, rel_pos_embeddings ) # [B, #heads, T, T'] or [B, #heads, T, T+T'+1] + attn_bd = self._rel_shift_bhij(attn_bd, k_len=time_dim_size) # [B, #heads, T, T'] - # We use attn_mask to add BD matrix to attention scores. - # - # Inside torch's SDPA the mask is added after regular scaling, so to get correct - # results, we need to apply the scaling here before passing to SDPA. - # - # See for reference: - # https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html - attn_bd_mask = attn_bd + mask - scale = math.sqrt(1.0 / float(self.embed_dim_per_head)) - attn_bd_mask_scaled = attn_bd_mask * scale + attn = attn_ac + attn_bd + mask # [B, #heads, T, T'] + attn_scaled = attn * (math.sqrt(1.0 / float(self.embed_dim_per_head))) # [B, #heads, T, T'] - v = value_seq.view(batch_dim_size, -1, self.num_heads, self.embed_dim_per_head) # [B, T, H, F'] + # softmax and dropout + attn_output_weights = self.att_weights_dropout(F.softmax(attn_scaled, dim=-1)) # [B, #heads, T, T'] - # Use torch's SDPA for efficiency. - # - # The attention matrices a and c are computed inside torch's sdpa. - attn_output = F.scaled_dot_product_attention( - q_with_bias_u.transpose(-3, -2), # [B, #heads, T, F'] - k.transpose(-3, -2), # [B, #heads, T', F'] - v.transpose(-3, -2), # [B, #heads, T, F'] - attn_mask=attn_bd_mask_scaled, # [B, #heads, T, T'] - dropout_p=self.att_weights_dropout.p if self.training else 0.0, - scale=scale, - ) # [B, #heads, T, F'] - attn_output = attn_output.transpose(-3, -2).flatten(-2) # [B, T, F'] - assert attn_output.shape[-1] == self.embed_dim + # sequence of weighted sums over value sequence + v = value_seq.view(batch_dim_size, -1, self.num_heads, self.embed_dim_per_head) # [B, T, H, F'] + attn_output = torch.einsum("bhij, bjhf -> bihf", attn_output_weights, v).reshape( + batch_dim_size, -1, self.embed_dim + ) output_tensor = self.out_proj(attn_output) + output_tensor = self.dropout(output_tensor) return output_tensor # [B,T,F] From 452daad3830dfd60d6b5eab6ff3d6ba6b3a9ceb8 Mon Sep 17 00:00:00 2001 From: Jingjing Xu Date: Fri, 17 Apr 2026 10:52:26 -0400 Subject: [PATCH 3/4] limit rel-pos change to learnable indexing --- i6_models/parts/conformer/mhsa_rel_pos.py | 59 +++++++++++++++-------- 1 file changed, 38 insertions(+), 21 deletions(-) diff --git a/i6_models/parts/conformer/mhsa_rel_pos.py b/i6_models/parts/conformer/mhsa_rel_pos.py index 820a5fdf..e5d6f16e 100644 --- a/i6_models/parts/conformer/mhsa_rel_pos.py +++ b/i6_models/parts/conformer/mhsa_rel_pos.py @@ -11,8 +11,8 @@ import torch.nn.functional as F from i6_models.config import ModelConfiguration -from i6_models.util import compat from i6_models.parts.dropout import BroadcastDropout +from i6_models.util import compat @dataclass @@ -172,10 +172,16 @@ def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> to 1, 2 * time_dim_size - 1, self.pos_emb_dim ) # [1, T+T'-1, pos_emb_dim] else: - rel_pos_embeddings = self._sinusoidal_pe( - torch.arange(time_dim_size - 1, -time_dim_size, -1, device=input_tensor.device, dtype=torch.float32), - self.pos_emb_dim, - ).view(1, 2 * time_dim_size - 1, self.pos_emb_dim) # [1, T+T'-1, pos_emb_dim] + rel_pos_embeddings = ( + self._sinusoidal_pe( + torch.arange( + time_dim_size - 1, -time_dim_size, -1, device=input_tensor.device, dtype=torch.float32 + ), + self.pos_emb_dim, + ) + .to(input_tensor.dtype) + .view(1, 2 * time_dim_size - 1, self.pos_emb_dim) + ) # [1, T+T'-1, pos_emb_dim] # dropout relative positional embeddings rel_pos_embeddings = self.pos_emb_dropout( @@ -194,30 +200,41 @@ def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> to q_with_bias_u = q + self.pos_bias_u if self.with_pos_bias else q # [B, T, #heads, F'] q_with_bias_v = q + self.pos_bias_v if self.with_pos_bias else q - # attention matrix a and c - attn_ac = torch.einsum("bihf, bjhf -> bhij", q_with_bias_u, k) # [B, #heads, T, T'] - # attention matrix b and d attn_bd = torch.einsum( - "bihf, ijhf -> bhij", q_with_bias_v, rel_pos_embeddings + "bihf, ijhf -> bhij", + q_with_bias_v, + rel_pos_embeddings.to(device=q_with_bias_v.device, dtype=q_with_bias_v.dtype), ) # [B, #heads, T, T'] or [B, #heads, T, T+T'+1] - attn_bd = self._rel_shift_bhij(attn_bd, k_len=time_dim_size) # [B, #heads, T, T'] + # We use attn_mask to add BD matrix to attention scores. + # + # Inside torch's SDPA the mask is added after regular scaling, so to get correct + # results, we need to apply the scaling here before passing to SDPA. + # + # See for reference: + # https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html + attn_bd_mask = attn_bd + mask + scale = math.sqrt(1.0 / float(self.embed_dim_per_head)) + attn_bd_mask_scaled = attn_bd_mask * scale - attn = attn_ac + attn_bd + mask # [B, #heads, T, T'] - attn_scaled = attn * (math.sqrt(1.0 / float(self.embed_dim_per_head))) # [B, #heads, T, T'] - - # softmax and dropout - attn_output_weights = self.att_weights_dropout(F.softmax(attn_scaled, dim=-1)) # [B, #heads, T, T'] - - # sequence of weighted sums over value sequence v = value_seq.view(batch_dim_size, -1, self.num_heads, self.embed_dim_per_head) # [B, T, H, F'] - attn_output = torch.einsum("bhij, bjhf -> bihf", attn_output_weights, v).reshape( - batch_dim_size, -1, self.embed_dim - ) - output_tensor = self.out_proj(attn_output) + # Use torch's SDPA for efficiency. + # + # The attention matrices a and c are computed inside torch's sdpa. + attn_output = F.scaled_dot_product_attention( + q_with_bias_u.transpose(-3, -2), # [B, #heads, T, F'] + k.transpose(-3, -2), # [B, #heads, T', F'] + v.transpose(-3, -2), # [B, #heads, T, F'] + attn_mask=attn_bd_mask_scaled, # [B, #heads, T, T'] + dropout_p=self.att_weights_dropout.p if self.training else 0.0, + scale=scale, + ) # [B, #heads, T, F'] + attn_output = attn_output.transpose(-3, -2).flatten(-2) # [B, T, F'] + assert attn_output.shape[-1] == self.embed_dim + output_tensor = self.out_proj(attn_output) output_tensor = self.dropout(output_tensor) return output_tensor # [B,T,F] From 0fd445ad49f5281af9dbc07e1fb51bd8cd80281b Mon Sep 17 00:00:00 2001 From: Jingjing Xu Date: Fri, 17 Apr 2026 11:22:54 -0400 Subject: [PATCH 4/4] update --- i6_models/parts/conformer/mhsa_rel_pos.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/i6_models/parts/conformer/mhsa_rel_pos.py b/i6_models/parts/conformer/mhsa_rel_pos.py index e5d6f16e..821b1c10 100644 --- a/i6_models/parts/conformer/mhsa_rel_pos.py +++ b/i6_models/parts/conformer/mhsa_rel_pos.py @@ -164,10 +164,9 @@ def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> to if self.learnable_pos_emb: # 1D optimization: 2T-1 unique relative positions instead of T×T distance matrix. - # Descending order matches sinusoidal branch convention for _rel_shift_bhij. - rel_pos = torch.arange(time_dim_size - 1, -time_dim_size, -1, device=input_tensor.device) - # _rel_shift_bhij produces i-j convention; embedding table uses j-i, so negate - indices = torch.clamp(-rel_pos, -self.rel_pos_clip, self.rel_pos_clip) + self.rel_pos_clip + # Build [-(T-1), ..., -1, 0, 1, ..., T-1] directly. + rel_pos = torch.arange(-(time_dim_size - 1), time_dim_size, device=input_tensor.device) + indices = torch.clamp(rel_pos, -self.rel_pos_clip, self.rel_pos_clip) + self.rel_pos_clip rel_pos_embeddings = self.rel_pos_embeddings[indices].view( 1, 2 * time_dim_size - 1, self.pos_emb_dim ) # [1, T+T'-1, pos_emb_dim]