Predicting High Uncertainty Events to Train Working Memory ( выжимка )
Artyom Sorokin | Dec 2021
Theory
Memory's Objective
Мы бы хотели чтобы наша память максимизировала следующую сумму:
argmaxθ∑t=0TI(ft;mt=gθ(c0,...,ct)∣сt)
argmax_{\theta} \sum_{t=0}^{T} I(f_{t} ; m_t = \textcolor{blue}{g_{\theta} (c_0, ..., c_t)} | с_t)
argmaxθ∑t=0TI(ft;mtθ∣сt)
argmax_{\theta} \sum_{t=0}^{T} I(f_{t} ; m^{\textcolor{blue}{\theta}}_t| с_t)
Для простоты буду писать так:
Memory's Objective
Проблема: учить память предсказывать будущее ft на шаге t может быть уже поздно:
argmaxθ∑t=0TI(ft;mtθ∣сt)
argmax_{\theta} \sum_{t=0}^{T} I(f_{t} ; m^{\textcolor{blue}{\theta}}_t| с_t)
Цель обучения памяти:
информация с шага t−k уже потеряна на шаге t
c0=[s0,a0]
c_0 = [s_0, a_0]
ct=[st,at]
c_t = [s_t, a_t]
ft
f_t
Memory's Objective
Чтобы выучится не выкидывать старую информацию придется на каждом шаге оптимизировать память относительно всех будущих шагов:
O(T2) по времени
Идея 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
насколько память может быть важна для предсказания ft
\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θ∗, тогда:
мелкая игнорим
оцениваем детектором
Общая схема обучения:
- Учим детектор dψ предсказывать ft на каждом шаге на основе ct. Важно уметь давать оценку неопределенности предсказаний детектора H^ψ(ft∣сt).
- Учим память gθ для каждого шага t предсказывать будущие события Ut, где память может быть наиболее важна:
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})\,,
Ut это набор шагов из эпизода для, которых детектор дал наибольшую оценку H^ψ(fi∣сi); ∣U∣≪T.
При обучении памяти мы используем информацию из будущего, которой не будем владеть во время её применения, поэтому нужна отдельная сетка для обьединения шагов k и t: предиктор
Optimization
Учим детектор
Нужно чтобы детектор dψ умел оценивать энтропию H(ft∣ct,mt):
H(ft∣ct)=Ep(ft∣ct)[−logp(ft∣ct)]
H(f_t|c_t) = \mathbb{E}_{p(f_t|c_t)}[\textcolor{blue}{-log\, p(f_t|c_t)}]
DKL(p(ft∣ct)∣∣dψ(ft∣∣ct))=CE(p(ft∣ct),dψ(ft∣ct))−H(p(ft∣ct))
\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))
Чтобы оценить −logp(ft∣ct), достаточно использовать Cross-Entropy Loss:
Мы не можем повлиять на H(ft∣ct), поэтому минимизируя CE loss мы минимизируем DKL между нашей моделькой и настроящим распределением.
−logdψ(ft∣st)
\textcolor{green}{- log \, d_{\psi}(f_t|s_t)}
если не сработает, будем искать более честную оценку энтропии
Наша оценка энтропии для бедных:
оценка энтропии по одному сэмплу или удивление
(1)
Учим память+предиктор
Память gθ и предиктор qϕ должны максимизировать взаимную информацию :
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(fk∣mtθ,ck):
\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))
Т.к.
DKL≥0
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θ и предиктор qϕ должны максимизировать взаимную информацию :
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, минимизируем её и разом обновляем θ и ϕ
Итог: dψ, gθ, qϕ учим предсказывать будущее ft при помощи NLL loss
Учим память+предиктор
На всякий случай для проверки, можно прочитать следующие 2 статьи из которых взято доказательство для функции memory+predictor:
Predicting High Uncertainty Events to Train Working Memory ( выжимка ) Artyom Sorokin | Dec 2021
MemUP Выжимка
By supergriver
MemUP Выжимка
- 494
