Efficient Reinforcement Learning by Reducing Forgetting with Elephant Activation Functions

Taylor Expansion

First-order Taylor expansion

$$f(x)=f(a)+f^{\prime }(a)(x-a)+O\left((x-a{)}^2\right)$$

Second-order Taylor expansion

$$f(x)=f(a)+f^{\prime }(a)(x-a)+\frac{1}{2}{f}^{\prime \prime }(a)(x-a{)}^2+O\left((x-a{)}^3\right)$$

Catastrophic Forgetting via Training Dynamics

Let a scalar-valued function $f_w(x)$ be represented by a neural network, parameterized by $w$, with input $x$. $F(x)$ is the true value function and $L(f,F,x)$ is the loss function.

$$w^{\prime }=w+{\Delta }_w$$

If using SGD,

$$w^{\prime }=w-\alpha {\nabla }_{w}L(f,F,x_t)$$ $${\Delta }_w=w^{\prime }-w=\alpha {\nabla }_{w}L(f,F,x_t)=-\alpha {\nabla }_{f}L(f,F,x_t){\nabla }_w{f}_w(x_t)$$

and the paper gives Equation 1 using the Taylor expansion:

$$f_w^{\prime }(x)-f_w(x)-\alpha {\nabla }_{f}L(f,F,x_t)⟨{\nabla }_w{f}_w(x_t),{\nabla }_w{f}_w(x_t)⟩+O({\Delta }_w^2)$$

I might want to show the derivation. Let

$$g(w):=f_w(x)$$

which expands to

$$g(w+{\Delta }_w)=g(w)+{\nabla }_{w}g(w{)}^{\top }{\Delta }_w=O(||{\Delta }_w||{)}^2$$

Now substitute back:

$$g(w+{\Delta }_w)=f_{w+{\Delta }_w}(x)=f_{w^{\prime }}(x)$$

and

$${\nabla }_{w}g(w{)}^{\top }{\Delta }_w={\nabla }_w{f}_w(x{)}^{\top }{\Delta }_w$$

Looking back at the original $g(w+{\Delta }_w)$, we now have

$$f_{w^{\prime }}(x)=f_w(x)+{\nabla }_w{f}_w(x{)}^{\top }{\Delta }_w+O(||{\Delta }_w|{|}^2)$$

and the neural tangent kernel (NTK) is $⟨{\nabla }_w{f}_w(x),{\nabla }_w{f}_w(x_t)⟩$.

I personally write it as $K_w(x,x_t)=⟨{\nabla }_w{f}_w(x),{\nabla }_w{f}_w(x_t)⟩$.

The paper argues for three properties:

• Property 1 (error correction). For $x=x_t$, $⟨{\nabla }_w{f}_w(x),{\nabla }_w{f}_w(x_t)⟩\ne 0$.
• Property 2 (zero forgetting). For $x\ne x_t$, $⟨{\nabla }_w{f}_w(x),{\nabla }_w{f}_w(x_t)⟩=0$.
• Property 3 (mild forgetting). $⟨{\nabla }_w{f}_w(x),{\nabla }_w{f}_w(x_t)⟩\approx 0$ for $x$ that is dissimilar to $x_t$.

Property 1 is easy to see: if $x=x_t$ and there is a prediction error, the kernel should not be zero, otherwise no update is being done.

The same goes for Property 2: if $x\ne x_t$, we should not update any of the parameters that affect $x$, so the kernel has to be zero for absolutely zero forgetting.

Property 3 relaxes this: the kernel value is mild, approximately zero. There will still be interference and forgetting, but only mildly.

Success and Failure of Sparse Representation

The paper states that sparse representations only reduce forgetting in linear function approximation, but not in nonlinear function approximation.

Linear function approximation, e.g. $f_w(w)=w^{\top }\varphi (x)$, where $phi(x)$ is fixed and only $w$ is trained. The prediction is a weighted sum of fixed features:

$$\varphi (x)=\left(\begin{array}{c}tanh(x)\\ tanh(x)\end{array}\right)$$

,

$$f_w(x)=w_1{\varphi }_1(x)=w_2{\varphi }_2(x).$$

Nonlinear function approximation

$$f_{u,v,b}(x)=utanh(vx+b)$$

where $u,v,b$ are trainable. The activation tanh has no parameters of its own, but the trainable $v,b$ are inside it, so when $v$ or $b$ changes, the feature itself changes. The features are learned.

In the linear case, it is easy to see that ${\nabla }_w{f}_w(x)=\varphi (x)$. Together with Equation 1,

$$f_w^{\prime }(x)-f_w(x)-\alpha {\nabla }_{f}L(f,F,x_t)\varphi (x{)}^{\top }\varphi (x)$$

Property 1 holds because $||\varphi (x_t)|{|}_2^2\ne 0$ when $x=x_t$. Property 3 holds because, when $x=x_t$, $⟨\varphi (x),\varphi (x_t)⟩\approx 0$. This is why sparse representation mitigates catastrophic forgetting in linear function approximation.

In the nonlinear case, using ${\varphi }_{\theta }(x)=\sigma (Vx+b)$, we have

$$⟨{\nabla }_w{f}_w(x),{\nabla }_w{f}_w(x_t)⟩={\varphi }_{\theta }^{\top }{\varphi }_{\theta }+u^{\top }u(x^{\top }{x}_t+1)+{{\varphi }_{\theta }^{\prime }}^{\top }{\varphi }_{\theta }^{\prime }$$

Property 3 does not hold here, even if ${\varphi }_{\theta }^{\top }{\varphi }_{\theta }\approx 0$, because it is not guaranteed that $u^{\top }u(x^{\top }{x}_t+1)+{{\varphi }_{\theta }^{\prime }}^{\top }{\varphi }_{\theta }\approx 0$.

Sparsity

$$A_{ϵ,C}=\left\{x\in [-C,C]:|\sigma (x)|\le ϵ\right\}.$$ $$S_{ϵ,C}(\sigma )=\frac{\left|\left\{x\mid |\sigma (x)|\le ϵ,x\in [-C,C]\right\}\right|}{\left|\left\{x\mid x\in [-C,C]\right\}\right|}=\frac{\left|\left\{x\mid |\sigma (x)|\le ϵ,x\in [-C,C]\right\}\right|}{2C}.$$

This is basically just

$$S_{ϵ,C}(\sigma )=\frac{\text{length of }{A}_{ϵ,C}}{\text{length of }\left[-C,C\right]}=\frac{\text{length of inputs producing near-zero outputs}}{2C}.$$

We can then define the sparsity $S(\sigma )$ as

$$S(\sigma )=\underset{ϵ\to 0^{+}}{\lim }\underset{C\to \infty }{\lim }{S}_{ϵ,C}(\sigma ).$$

and we call $\sigma$ a sparse function when $S(\sigma )=1$.

For example, tanh'(x) is a sparse function, while tanh(x) itself is not.

Elephant Activation Functions

$$Elephant(x)=\frac{h}{1+|\frac{x}{a}{|}^d}$$

where $a$ controls the width of the function and hence the sparsity of the function itself, $h$ controls the height of the function, and $d$ controls the sparsity of the gradients of the elephant function.