Predicting High Uncertainty Events to Train Working Memory ( выжимка )

Artyom Sorokin | Dec 2021

Theory

Memory's Objective

Мы бы хотели чтобы наша память максимизировала следующую сумму: 

argmax_{\theta} \sum_{t=0}^{T} I(f_{t} ; m_t = \textcolor{blue}{g_{\theta} (c_0, ..., c_t)} | с_t)
argmax_{\theta} \sum_{t=0}^{T} I(f_{t} ; m^{\textcolor{blue}{\theta}}_t| с_t)

Для простоты буду писать так: 

Memory's Objective

Проблема: учить память предсказывать будущее \(f_t\) на шаге t может быть уже поздно: 

argmax_{\theta} \sum_{t=0}^{T} I(f_{t} ; m^{\textcolor{blue}{\theta}}_t| с_t)

Цель обучения памяти: 

информация с шага \(t-k\) уже потеряна на шаге \(t\)

c_0 = [s_0, a_0]
c_t = [s_t, a_t]
f_t

Memory's Objective

Чтобы выучится не выкидывать старую информацию придется на каждом шаге оптимизировать память относительно всех будущих шагов:  

\(O(T^2)\) по времени

Идея MemUP:
вместо того, чтобы оптимизировать  вторую сумму полностью, выберем шаги когда обучение памяти  может дать наибольший вклад в предсказание будущего!

\sum_{t=0}^{T} \textcolor{blue}{\sum_{i=t}^{T}} I(f_\textcolor{blue}{i} ; m_\textcolor{black}{t} | с_\textcolor{blue}{i})

все еще нужно обрабатвать всю последовательность длинны \(T\)

Если поменять порядок сумм то получится 1 в 1 как учатся трансформер в RL

Finding Important Moments

насколько память может быть важна  для предсказания \(f_t\)

\textcolor{black}{\sum_{t=0}^{T}} I(f_\textcolor{black}{t} ; m^\textcolor{black}{*}_t | c_\textcolor{black}{t})
= \textcolor{black}{\sum_{t=0}^{T}} [H(f_\textcolor{black}{t}| c_\textcolor{black}{t}) - H(f_\textcolor{black}{t}|m^*, c_\textcolor{black}{t})]

Найдем моменты в будущем, когда выучивание памяти потенциально может принести максимальную пользу

Вообразим, что у нас есть идеальная память \(m^{\theta^*}\), тогда: 

мелкая игнорим

оцениваем детектором

Общая схема обучения:

  • Учим детектор \(d_{\psi}\) предсказывать \(f_t\) на каждом шаге на основе \(c_t\). Важно уметь давать оценку  неопределенности предсказаний детектора \(\hat{H}_{\psi}(f_t| с_t)\).

 

  • Учим память \(g_{\theta}\) для каждого шага t предсказывать будущие события \(U_t\), где память может быть наиболее важна:​
argmax_{\theta} \sum_{t=0}^{T} \textcolor{green}{\sum_{k \in U_t}} I(f_\textcolor{green}{k} ; m^{\theta}_\textcolor{black}{t} | с_\textcolor{green}{k})\,,

\(U_t\) это набор шагов из эпизода для, которых детектор дал наибольшую оценку \(\hat{H}_{\psi}(f_i| с_i)\); \(|U| \ll T\).

При обучении памяти мы используем информацию из будущего, которой не будем владеть во время её применения, поэтому нужна отдельная сетка для обьединения шагов k и t: предиктор

Общая схема обучения:

argmax_{\theta} \sum_{t=0}^{T} \textcolor{black}{\sum_{k \in U_t}} I(f_\textcolor{black}{k} ; m^{\theta}_\textcolor{black}{t} | с_\textcolor{black}{k})\,,
\theta^* = argmax_{\theta} \sum_{t=0}^{T} \textcolor{black}{\sum_{k \in U_t}} I(f_\textcolor{black}{k} ; m_{t} = g_{\theta}(c_0,..., c_t) | с_\textcolor{black}{k})\,,
I(f_\textcolor{black}{k} ; m^\textcolor{black}{\theta}_t | с_\textcolor{black}{k}) = H(f_k|c_k) - H(f_k|m^\textcolor{black}{\theta}_t, c_k)
\ge H(f_k|c_k) - CE(p(f_k|m^\textcolor{black}{\theta}_t, c_k), \textcolor{black}{q_{\phi}}(f_k|m^\textcolor{black}{\theta}_t, c_k))
\theta^*, \phi^* = argmin_{\theta, \phi}\, CE(p(f_k|m^\textcolor{black}{\theta}_t, c_k), \textcolor{black}{q_{\phi}}(f_k|m^\textcolor{black}{\theta}_t, c_k)) = argmin_{\theta, \phi}\, \mathbf{E}_{p(f_k, m^\textcolor{black}{\theta}_t|c_k)}[ -log\, q_\textcolor{black}{\phi}(f_k| m^\textcolor{black}{\theta}_t, c_t)]
\psi^* = argmin_{\psi}\, CE(p(f_k|c_k), d_{\psi}(f_k| c_k))
I(f_\textcolor{black}{t} ; \hat{m}_t | c_\textcolor{black}{t})
= H(f_\textcolor{black}{t}| c_\textcolor{black}{t}) - H(f_\textcolor{black}{t}|\hat{m_t}, c_\textcolor{black}{t})
\textbf{D}_{KL} (p(f_t|c_t) || \textcolor{black}{d_{\psi}(f_t| c_t)}) = \textbf{CE}(p(f_t|c_t), \textcolor{black}{d_{\psi}(f_t|c_t)}) - H(p(f_t|c_t))

Optimization

Учим детектор

Нужно чтобы детектор \(d_{\psi}\) умел оценивать энтропию \(H(f_t|c_t)\):

H(f_t|c_t) = \mathbb{E}_{p(f_t|c_t)}[\textcolor{blue}{-log\, p(f_t|c_t)}]
\textbf{D}_{KL} (p(f_t|c_t) || \textcolor{green}{d_{\psi}(f_t|| c_t)}) = \textbf{CE}(p(f_t|c_t), \textcolor{green}{d_{\psi}(f_t|c_t)}) - H(p(f_t|c_t))

Чтобы оценить \(-log\, p(f_t|c_t)\), достаточно использовать Cross-Entropy Loss: 

Мы не можем повлиять на \(H(f_t|c_t)\), поэтому минимизируя CE loss мы минимизируем \(D_{KL}\) между нашей моделькой и настроящим распределением. 

\textcolor{green}{- log \, d_{\psi}(f_t|c_t)}

если не сработает, будем искать более честную оценку энтропии

Наша оценка энтропии для бедных:

То что CE минимизурует \(D_{KL}\), результат супер известный, можно даже не расписывать в статье.

(1)

оценка энтропии по одному сэмплу или удивление

Учим детектор

Нужно чтобы детектор \(d_{\psi}\) умел оценивать энтропию \(H(f_t|c_t, m_t)\):

H(f_t|c_t) = \mathbb{E}_{p(f_t|c_t)}[\textcolor{blue}{-log\, p(f_t|c_t)}]
\textbf{D}_{KL} (p(f_t|c_t) || \textcolor{green}{d_{\psi}(f_t|| c_t)}) = \textbf{CE}(p(f_t|c_t), \textcolor{green}{d_{\psi}(f_t|c_t)}) - H(p(f_t|c_t))

Чтобы оценить \(-log\, p(f_t|c_t)\), достаточно использовать Cross-Entropy Loss: 

Мы не можем повлиять на \(H(f_t|c_t)\), поэтому минимизируя CE loss мы минимизируем \(D_{KL}\) между нашей моделькой и настроящим распределением. 

\textcolor{green}{- log \, d_{\psi}(f_t|s_t)}

если не сработает, будем искать более честную оценку энтропии

Наша оценка энтропии для бедных:

оценка энтропии по одному сэмплу или удивление

(1)

Учим память+предиктор

Память \(g_{\theta}\) и предиктор \(q_{\phi}\)  должны максимизировать взаимную информацию :

Barber, Agakov (2004) доказали Lower Bound для взаимной информации

(перепишем для нашего случая):

\sum_{t=0}^{T} \textcolor{black}{\sum_{k \in U_t}} I(f_\textcolor{black}{k} ; m_t = \textcolor{green}{ g_{\theta}(h_{1:t})} | с_\textcolor{black}{k})\, = \sum_{t=0}^{T} \textcolor{black}{\sum_{k \in U_t}} I(f_\textcolor{black}{k} ; m^\textcolor{green}{\theta}_t | с_\textcolor{black}{k})
I(f_\textcolor{black}{k} ; m^\textcolor{green}{\theta}_t | с_\textcolor{black}{k}) = H(f_k|c_k) - H(f_k|m^\textcolor{green}{\theta}_t, c_k)

вспоминаем ур.1 между \(CE\) и \(KL\) только для распределения \(p(f_k|m^{\theta}_t, c_k)\):

\textbf{D}_{KL} (p(f_k|m^\textcolor{green}{\theta}_t, c_k)|| \textcolor{blue}{q_{\phi}}(f_k|m^\textcolor{green}{\theta}_t, c_k)) = \textbf{CE}(p(f_k|m^\textcolor{green}{\theta}_t, c_k), \textcolor{blue}{q_{\phi}}(f_k|m^\textcolor{green}{\theta}_t, c_k)) - H(p(f_k|m^\textcolor{green}{\theta}_t, c_k))

Т.к.

D_{KL} \ge 0

, значит

CE(p_\textcolor{black}{\theta}, q_{\phi,\theta}) \ge H(p_\theta)

 получается:  

(2)

H(f_k|c_k) - H(f_k|m^\textcolor{green}{\theta}_t, c_k)
\ge H(f_k|c_k) - CE(p(f_k|m^\textcolor{green}{\theta}_t, c_k), \textcolor{blue}{q_{\phi}}(f_k|m^\textcolor{green}{\theta}_t, c_k))

Учим память+предиктор

Память \(g_{\theta}\) и предиктор \(q_{\phi}\)  должны максимизировать взаимную информацию :

Barber, Agakov (2004) доказали Lower Bound для взаимной информации

(перепишем для нашего случая):

\sum_{t=0}^{T} \textcolor{black}{\sum_{k \in U_t}} I(f_\textcolor{black}{k} ; m_t = \textcolor{green}{ g_{\theta}(h_{1:t})} | с_\textcolor{black}{k})\, = \sum_{t=0}^{T} \textcolor{black}{\sum_{k \in U_t}} I(f_\textcolor{black}{k} ; m^\textcolor{green}{\theta}_t | с_\textcolor{black}{k})
H(f_k|c_k) - H(f_k|m^\textcolor{green}{\theta}_t, c_k)
\ge H(f_k|c_k) - CE(p(f_k|m^\textcolor{green}{\theta}_t, c_k), \textcolor{blue}{q_{\phi}}(f_k|m^\textcolor{green}{\theta}_t, c_k))
\ge H(f_k|c_k) - \mathbb{E}_{p(f_k, m^\textcolor{green}{\theta}_t|c_k)}[ -log\, q_\textcolor{blue}{\phi}(f_k| m^\textcolor{green}{\theta}_t, c_t)]
I(f_\textcolor{black}{k} ; m^\textcolor{green}{\theta}_t | с_\textcolor{black}{k})

от нас не зависит, игнор

NLL loss, минимизируем её  и разом обновляем \(\theta\) и \(\phi\)

Итог: \(d_\psi\), \(g_\theta\), \(q_\phi\)  учим предсказывать будущее \(f_t\) при помощи NLL loss

Учим память+предиктор

На всякий случай для проверки, можно прочитать следующие 2 статьи из которых взято доказательство для функции memory+predictor: 

Architecture

Architecture

 

  • TrajGenerator создаем последовательности
    • TargetCreator: может потрбоваться создавать цели предсказания (RL, self-supervision) 
  • UncertaintyDetector: оцениваем неопределенность шагов в траектории
    • СontextEncoder: нужен здесь и для предиктора
  • EventSelector: Чтобы по неопределенности выбрать события для предсказаний памяти
  • EpisodeBuffer: Выбирает батчи эпизодов, запускает на них UncertaintyDetector, EventSelector
    • ​Возможно нужен отдельный класс чтобы все собрать в батч 
  • MemoryModule: учится хранить информацию
  • Predictor: позволяет учить памят на событиях из будущего. 

 

Чтобы учить память используем всё. Чтобы учить детектор достаточно только первых 4ех

 

ContextEncoder

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

* Применяется перед использованием Detector'а и Predictor'а (должен быть одинаковый для этой пары)

Context

TrajGenerator

Какие варианты могут быть:

  • Просто создавать, как copy-task: CopyGenerator(<аргументы задачи>)
  • Загружать из датасета: DatasetLoader(<путь и тд.>)
  • RL версия: RLGenerator( env, policy, etc.)
    • SequentialRLGenerator(env, policy, etc.)
    • ParallelRLGenerator(env, policy, MultiprocessingQueue?)

 

Поля:
* get_trajs(n?)

* TrajIterator