File size: 907 Bytes
d9d4242 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | """Layer-wise attention pooling over transformer hidden states."""
from __future__ import annotations
import torch
import torch.nn as nn
class LayerwiseAttention(nn.Module):
def __init__(self, num_layers: int) -> None:
super().__init__()
self.num_layers = num_layers
self.weights = nn.Parameter(torch.zeros(num_layers))
self.gamma = nn.Parameter(torch.ones(1))
def forward(
self, hidden_states: tuple[torch.Tensor, ...] | list[torch.Tensor]
) -> torch.Tensor:
if len(hidden_states) != self.num_layers:
raise ValueError(
f"expected {self.num_layers} layers, got {len(hidden_states)}"
)
w = torch.softmax(self.weights, dim=0)
stacked = torch.stack(list(hidden_states), dim=0) # [L,B,T,H]
combined = (w.view(-1, 1, 1, 1) * stacked).sum(dim=0)
return self.gamma * combined
|