| """Simple learned-attention pooling. | |
| STATUS (as of v9, 2026-05-27): NOT integrated into the current model. | |
| Kept as a v10+ ablation baseline against the tree aggregator. Tests if | |
| the structural prior in TreeAggregator buys anything over a vanilla | |
| single-head learned-attention pooling. | |
| Has tests in tests/test_attention_pooling.py. | |
| """ | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| class AttentionPooling(nn.Module): | |
| def __init__(self, d_model: int) -> None: | |
| super().__init__() | |
| self.d_model = d_model | |
| self.W_doc = nn.Linear(in_features=d_model, out_features=1) # (d_model,1) | |
| def forward( | |
| self, | |
| h: torch.Tensor, # Shape(Batch, Seq_Len, d_model) | |
| ): | |
| scores: torch.Tensor = self.W_doc(h) # (Batch, Seq_Len, 1) | |
| scores = scores / math.sqrt(self.d_model) | |
| scores = scores.view(scores.shape[0], 1, scores.shape[1]) # (Batch, 1, Seq_Len) | |
| weights: torch.Tensor = scores.softmax(dim=-1) # (Batch,1, Seq_Len) | |
| return ( | |
| (weights @ h).squeeze(1), | |
| weights, | |
| ) # returning weights for testing/visualization (Batch,d_model), (Batch,1,Seq_Len) | |