You also need GatedMaxPool

[Mihaylo]

Я тут размышляю над причиной, почему maxpool неплохо работает вместо трансформеров.

Трансформерный блок имеет сигнатуру (B, T, D) -> (B, T, D), где B - размер батча, T - длина последовательности, D - длина вектора данных. Искусственный интеллект сразу разобрался и посоветовал замену с шагом S = 1 и левым паддингом P = W - 1, где W - размер окна. В таком случае GateTopkPool подходит по сигнатуре вместо трансформерного блока. В процессе исследования я также установил, что вариант Topk = 1 (GatedMaxPool вместо тормознутого GatedTopkPool), W = 2 даёт лучшие результаты. В итоге рекомендуемые параметры: GatedMaxPool, W = 2, S = 1, P = [1, 0] - левый паддинг есть, правого нет. Также присутствует линейный слой с активацией с матрицей [D, D] - такой слой не изменяет размерность (B, T, D), просто он умный, с весами.

Я ведь пробовал заменить трансформеры на простые линейные слои с активацией и с матрицей [D, D], без maxpool. Эта штука не работает от слова совсем. Рецептивное поле равно 1. Слепота.

Линейные слои с активацией и с матрицей [D, D] способны управлять через maxpool передачей данных между соседними токенами, организуя рецептивное поле - вот почему работает GatedMaxPool.

P.S. Сравним с maxpool в классических нейросетях распознавания изображений. Основное отличие - это то, что у maxpool в распознавании изображений S = W. В изображениях maxpool выполняет настоящую роль пулинга - сжатия данных. Из-за того, что S больше единицы, рецептивное поле разрушается. Я бы попробовал ретроспективно вернуться к изображениям и попробовать S = 1 в maxpool и заменить при этом свёртку на линейные слои. Кажется это могло бы быть эффективной заменой свёрточных слоёв. Конечно же стандартный maxpool надо оставить, то есть скомбинировать два maxpool с S = 1 и S = W. При этом для первого варианта maxpool (который GatedMaxPool) можно поэкспериментировать W > 2, так как изображения более избыточны информативно. Можно выиграть в производительности вычислений.

[Mihaylo]

Есть очень правдоподобная гипотеза, почему GatedMaxPool сходится быстрее, чем Transformer - благодаря именно более узкому рецептивному полю - нейросеть не рассматривает связи за пределами рецептивного поля и это очень годный подход в случае анализа простых текстов. Однако на более сложных задачах Transformer превзойдёт GatedMaxPool в разы.

ИИ подсказывает решение в виде манипуляции с параметром dilation оконной функции. Надо сделать так и тогда рецептивное поле возрастёт экспоненциально ($R=2^N$) при сохранении количества слоёв:

        self.cmp1 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)
        self.cmp2 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=2)
        self.cmp3 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=4)
        self.cmp4 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=8)
        self.cmp5 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=16)
        self.cmp6 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=32)
        self.cmp7 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=64)
        self.cmp8 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=128)

Я правда удалил dilation из кода, когда выкладывал на форум, возвращаю:

import torch
import torch.nn as nn
from typing import Union, Tuple
import torch.nn.functional as F

class CausalGatedMaxPool(nn.Module):
    def __init__(
        self,
        d_model: int,
        window_size: int = 2,
        dilation: int = 1,
    ):
        super().__init__()

        if d_model < 1:
            raise ValueError("d_model must be >= 1")
        if window_size < 2:
            raise ValueError("window_size must be >= 2")
        if dilation < 1:
            raise ValueError("dilation must be >= 1")

        self.d_model = int(d_model)
        self.window_size = int(window_size)
        self.dilation = int(dilation)

        # Один обучаемый коэффициент на каждую 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.ReLU()
        
    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)}"
            )
    
        x = x * self.gate
    
        original_shape = x.shape
        T, D = original_shape[-2:]
        batch_shape = original_shape[:-2]
    
        # (..., T, D) -> (N, D, T)
        x = x.reshape(-1, T, D).transpose(1, 2)
    
        left_pad = (self.window_size - 1) * self.dilation
    
        # Важно: ручной asymmetric padding.
        # Аргумент padding у max_pool1d здесь не подходит:
        # он симметричный и ограничен.
        x = F.pad(x, (left_pad, 0), value=float("-inf"))
    
        # (N, D, T + left_pad) -> (N, D, T)
        pooled = F.max_pool1d(
            x,
            kernel_size=self.window_size,
            stride=1,
            padding=0,
            dilation=self.dilation,
        )
    
        # (N, D, T) -> (..., T, D)
        pooled = pooled.transpose(1, 2).reshape(*batch_shape, T, D)
    
        return self.act(self.proj(pooled))

Как работает dilation-трюк:



P.S. ИИшка на схеме почему-то слово dilated заменил на delayed.

[Mihaylo]

Есть вот такие теоретические изыскания:

1. Плотное (но узкое) перцептивное поле, назовём его равномерным, так как каждый входной токен в пределах перцептивного поля равнозначен и участвует в "голосовании" наравне с другими:

        self.cmp1 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)
        self.cmp2 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)
        self.cmp3 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)
        self.cmp4 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)

2. Сильно разреженное в прошлом (но широкое) рецептивное - экспоненциальное. Прошлые токены получают меньше внимания.

        self.cmp1 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)
        self.cmp2 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=2)
        self.cmp3 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=4)
        self.cmp4 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=8)

3. Промежуточный вариант - фибоначчиевый:

        self.cmp1 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)
        self.cmp2 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=2)
        self.cmp3 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=3)
        self.cmp4 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=5)

4. Глобальный вариант - трансформер. Тоже можно использовать.

5. Комбинированный вариант:

        self.cmp1 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)
        self.cmp2 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)
        self.cmp3 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=1)
        self.cmp4 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=2)
        self.cmp5 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=3)
        self.cmp6 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=4)
        self.cmp7 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=8)
        self.cmp8 = CausalGatedMaxPool(d_model=D_MODEL, window_size=2, dilation=16)
        self.transformer1 = Transformer()
        self.transformer2 = Transformer()

Все представленные варианты не имеют дыр в рецептивном поле. Самый разреженный вариант - экспоненциальный - имеет самое широкое рецептивное поле, но повышенную слепоту в прошлое.

Комбинируя слои GPT-нейросети, можно добиться наилучшего соответствия архитектуры и данных, и в итоге получить более лёгкие вычисления и одновременно более высокую сходимость.

GatedMaxPool изучает зависимости между соседними токенами, что соответствует слоговым знаниям, частям слова и далее более глубокие слои GatedMaxPool изучают устойчивые словосочетания, завершается всё трансформерами, которые способны находить логические связи в предложениях, абзацах, главах, разделах, включая начало и конец текста.