Ниже приведён код каузального варианта GatedMaxPool. Он может применяться для замены трансформерных блоков в различных GPT-нейросетях.
Код:
class CausalGatedMaxPool(nn.Module):
def __init__(
self,
d_model: int,
window_size: int = 2,
):
super().__init__()
if d_model < 1:
raise ValueError("d_model must be >= 1")
if window_size < 2:
raise ValueError("window_size must be >= 2")
self.d_model = int(d_model)
self.window_size = int(window_size)
# Один обучаемый коэффициент на каждую feature-компоненту.
# Инициализация около ±1 сохраняет и max-, и min-пути с самого начала.
self.gate = nn.Parameter(torch.empty(d_model))
nn.init.normal_(self.gate, mean=0.0, std=1.0)
self.proj = nn.Linear(d_model, d_model)
self.act = nn.GELU()
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
x: (..., T, d_model)
returns: (..., T, d_model)
"""
if x.ndim < 2:
raise ValueError("Expected x with shape (..., T, d_model)")
if x.size(-1) != self.d_model:
raise ValueError(
f"Expected last dim {self.d_model}, got {x.size(-1)}"
)
# (..., T, D): feature-wise статический гейт.
# Гейт применяем ПЕРЕД padding: отрицательный gate безопасен.
x = x * self.gate
# Реальная ширина сегмента до отбора прореженных позиций.
left_pad = (self.window_size - 1) * self.dilation
# Padding не должен стать максимумом ни при каких данных.
x = F.pad(
x,
pad=(0, 0, left_pad, 0),
mode="constant",
value=float("-inf"),
)
# (..., T, D, left_pad + 1)
segments = x.unfold(
dimension=-2,
size=left_pad + 1,
step=1,
)
# (..., T, D)
pooled = segments.amax(dim=-1)
return self.act(self.proj(pooled))
Типовое применение:
Код:
self.cgmp1 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2)
self.cgmp2 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2)
self.cgmp3 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2)
self.cgmp4 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2)
self.cgmp5 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2)
self.cgmp6 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2)
self.cgmp7 = CausalGatedMaxPool(d_model=D_MODEL, window_size=16)
self.cgmp8 = CausalGatedMaxPool(d_model=D_MODEL, window_size=16)
...
def forward(self, idx, targets=None):
#positions = torch.arange(time_steps, device=idx.device)
x = self.token_embedding(idx) # + self.position_embedding(positions)
#x = self.blocks(x)
x = x + self.cgmp1(x)
x = x + self.cgmp2(x)
x = x + self.cgmp3(x)
x = x + self.cgmp4(x)
x = x + self.cgmp5(x)
x = x + self.cgmp6(x)
x = x + self.cgmp7(x)
x = x + self.cgmp8(x)
...
По сути, почти обычный макспулинг в комбинации с линейными слоями и активациями обучаются как языковые модели без использования трансформеров.
Просто случайное открытие.Представленный слой сумел выучить шекспировский датасет и воспроизвести в его стиле следующее:
Цитата:
ROMEO:
I will affairs of the time modesty.
Clown:
The unhappy masters well, will not fear this true.
ROMEO:
And, and your late, nor with pardon are saw the
restite with the pardon fourth the wife,
Thou wilt to sits me from such every to the self;
And what is your words?
KING RICHARD III:
No; and down to thy beats.
ESCALUS:
And with me to make some to unship
bright it down to be prepared that should I saw mad;
The offence is famous subject my father,
Is so right for marriage of all be sitcuse noise.
Second Servingman:
I think you'll for your brother is
Stand with then do access that we Servingman:
This is merry be gone.
GLOUCESTER:
Had, thou art to take up, and yet lead me but made
And your skill hand shame hither speech
That I come in the my soul steel we and mine either
Your of lacks the first myself above takers upon.
GLOUCESTER:
An or the way, not shall be no bed
The sing to against thou seem, 'tis ere place,
Or prisoner and her tongue a soldier.
Somerset, which I lead not with sir, with me in any curse the truly to the tent.
GLOUCESTER:
Well so he might bed myselves one most thy first George!
CATES:
And I'll that see, and which man's rest
I can his policious both, to do such kiss me about.
Shepherish and straign living surping.
KING HENRY VI:
Have helm, knife: but even is a beward?
POLIXENES:
He wast threshes, beholds where is such of precious have let our brain,
Who say the true well a favouring from the gods?
DUKE VINCENTIO:
All to hence, most They with a pardon,
That made of me by with his modestrious daughter, came and her pass'd with the orators i' the matter than thou have being with night, here come rest.
The soul and death it with virgininuries! and my father
being to prove this plarus wouldst be,
Which this full of your friends of the course.
LADY ANNE:
Prove it, sir:
Every true his supposed with and seems.
JULIET:
O, sir, proffer's perjury, so me your princes,
So man should pride the good talk of the was a shame.
Provost:
Why I dare glass ca
![$$
x^{t, \ell + 1}_{\mu} = x^{t, \ell}_{\mu} +
{\sf GELU}
\left[
\sum_{\nu = 0}^{D - 1} M^{\ell}_{\mu \nu}
{\sf MAX}\left[
\left(x^{t - 0, \ell}_{\nu} \cdot g^{0, \ell}_{\nu} \right),
\,
\left(x^{t - 1, \ell}_{\nu} \cdot g^{1, \ell}_{\nu} \right),
\,
\ldots,
\left(x^{t - (W-1), \ell}_{\nu} \cdot g^{W-1, \ell}_{\nu} \right)
\right]
\right]
$$](https://dxdy.ru/math/890d31997f2124f9d9407f337b42965082.png "$$
x^{t, \ell + 1}_{\mu} = x^{t, \ell}_{\mu} +
{\sf GELU}
\left[
\sum_{\nu = 0}^{D - 1} M^{\ell}_{\mu \nu}
{\sf MAX}\left[
\left(x^{t - 0, \ell}_{\nu} \cdot g^{0, \ell}_{\nu} \right),
\,
\left(x^{t - 1, \ell}_{\nu} \cdot g^{1, \ell}_{\nu} \right),
\,
\ldots,
\left(x^{t - (W-1), \ell}_{\nu} \cdot g^{W-1, \ell}_{\nu} \right)
\right]
\right]
$$")
