Блог

Уровень AttentionQKV, реализованный Trax, выглядит следующим образом: AttentionQKV

 def AttentionQKV(d_feature, n_heads=1, dropout=0.0, mode='train'):
  """Returns a layer that maps (q, k, v, mask) to (activations, mask).
  See `Attention` above for further context/details.
  Args:
    d_feature: Depth/dimensionality of feature embedding.
    n_heads: Number of attention heads.
    dropout: Probababilistic rate for internal dropout applied to attention
        activations (based on query-key pairs) before dotting them with values.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
  """
  return cb.Serial(
      cb.Parallel(
          core.Dense(d_feature),
          core.Dense(d_feature),
          core.Dense(d_feature),
      ),
      PureAttention(  # pylint: disable=no-value-for-parameter
          n_heads=n_heads, dropout=dropout, mode=mode),
      core.Dense(d_feature),
  )
  

В частности, какова цель трех параллельных плотных слоев? Входными данными для этого слоя являются q, k, v, маска. Почему q, k, v проходят через плотный слой?

Этот фрагмент кода представляет собой реализацию уравнения в верхней части страницы 5 документа «Внимание — это все, что вам нужно«, в котором в 2017 году были представлены модели Transformer. Вычисление проиллюстрировано на рисунке 2 статьи:

Скрытые состояния проецируются на h головок внимания, которые параллельно выполняют масштабированное точечное произведение внимания. Проекция может быть интерпретирована как извлечение информации, которая имеет отношение к голове. Затем каждая головка выполняет вероятностный поиск на основе разных (изученных) критериев.