Масштабируемый, эффективный иерархический Softmax в Tensorflow?

Я заинтересован в реализации иерархической модели softmax, которая может обрабатывать большие словари, например, порядка 10M классов. Каков наилучший способ сделать это для того, чтобы быть масштабируемым до большого класса и эффективным? Например, по крайней мере одна статья показала, что HS может достичь ускорения ~ 25x для больших словаров при использовании двухуровневого дерева, где каждый node sqrt(N). Меня интересует также более общая версия для произвольного дерева глубины с произвольным коэффициентом ветвления.

Есть несколько вариантов, которые я вижу здесь:

1) Запустите tf.gather для каждой партии, где мы собираем индексы и разбиваем. Это создает проблемы с большими размерами партий и жировыми деревьями, где теперь коэффициенты дублируются, что приводит к ошибкам OOM.

2) Подобно # 1, мы могли бы использовать tf.embedding_lookup, который мог бы помочь с ошибками OOM, но теперь держит все на процессоре и немного замедляет работу.

3) Используйте tf.map_fn с помощью parallel_iterations=1, чтобы обрабатывать каждый образец отдельно и вернуться к использованию сбора. Это намного более масштабируемо, но на самом деле не приближается к 25-кратному ускорению из-за сериализации.

Есть ли лучший способ реализовать HS? Существуют ли разные способы для глубоких и узких или коротких и широких деревьев?

ОТВЕТЫ

Ответ 1

Вы упомянули, что хотите производительность GPU-класса:

но теперь держит все на процессоре и немного замедляет работу

и хотите использовать скрытый размер 300 единиц и словари на 10 М слов.

Это означает, что (при условии float32) вам понадобится 4 * 300 * 10M * 2 байта = 24 ГБ, чтобы сохранить параметры и градиент для выходного слоя.

Иерархический Softmax (HSM) не уменьшает требования к памяти - он просто ускоряет обучение.

Реально, вам понадобится намного больше памяти GPU, потому что вам также нужно будет хранить:

  • другие параметры и их градиенты

  • данные оптимизатора, например. скорости в обучении импульсам

  • активация и backpropagated временных данных

  • служебные данные, связанные с инфраструктурой

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

Однако у вас теперь есть другая проблема:

Чтобы сделать это конкретным, предположим, что у вас есть двухуровневый HSM с 3K классами, с 3K словами на класс (всего 9M слов). Вы распределяете классы 3K по 8 графическим процессорам, так что каждый хост 384 класса.

Что делать, если все целевые слова в пакете из одних и тех же классов 384, т.е. принадлежат одному и тому же GPU? Один GPU будет выполнять всю работу, в то время как остальные 7 ждут его.

Проблема в том, что даже если целевые слова в пакете принадлежат к различным графическим процессорам, вы все равно будете иметь такую ​​же производительность, как в худшем случае, если вы хотите сделать это вычисление в TensorFlow (это потому, что TensorFlow представляет собой структуру "указать-и-запустить" - вычислительный граф является одинаковым для наилучшего случая и наихудшего случая)

Каков наилучший способ сделать это для того, чтобы быть масштабируемым до большого класса и эффективным?

Вышеприведенная неэффективность модели parallelism (каждый GPU должен обрабатывать всю партию) предполагает, что нужно стараться держать все в одном месте.

Предположим, что вы либо реализуете все на хосте, либо на одном огромном графическом процессоре.

  • Если вы не моделируете последовательности или если вы есть, но для всей последовательности есть только один вывод, тогда накладные расходы памяти от копирования параметров, на которые вы ссылались, ничтожно по сравнению с требованиями к памяти описанных выше:

    400 == размер партии < количество классов == 3K

    В этом случае вы можете просто использовать gather или embedding_lookup (хотя копирование неэффективно)

  • Однако, если вы моделируете последовательности длины, скажем, 100, с выходом на каждом временном шаге, то копирование параметров становится большой проблемой.

    В этом случае, я думаю, вам нужно будет спуститься до С++/CUDA C и реализовать весь этот слой и его градиент как пользовательский op.