# Learning disentangled representations in a 'sphere-world' with continuous actions

## 1. Preamble

```python
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import random

import time
from IPython import display

try:
    import seaborn as sns
    sns.set()
    print("Set seaborn plotting defaults.")
except:
    print("seaborn not found.")
```

    Set seaborn plotting defaults.

```python
def mk_dir(export_dir, quite=False):
    if not os.path.exists(export_dir):
            try:
                os.makedirs(export_dir)
                print('created dir: ', export_dir)
            except OSError as exc: # Guard against race condition
                 if exc.errno != exc.errno.EEXIST:
                    raise
            except Exception:
                pass
    else:
        print('dir already exists: ', export_dir)

mk_dir("_data/")
```

    dir already exists:  _data/

## 2. Defining objects

### SphereWorld environment

We use a cubic 3D world in which a ball evolves on a sphere. Observations are a 3-dimensional tensor with values 0 except at the voxel occupied by the ball which has value 1. When the ball is not at the center of the voxel (which is most of the time) its density is distributed around neighboring voxels. Actions are continuous rotations around the three axes.

```python
class SphereWorld():

    class action_space():
        def __init__(self,n_actions):
            self.n = n_actions

        def sample(self, k=1):
            return torch.randint(0,self.n,(k,))

    class observation_space():
        def __init__(self,n_features):
            self.shape = [n_features]

    def __init__(self,dim=5,radius=3.5):

        self.size = dim

        self.radius = radius
        self.action_space = self.action_space(6)
        self.observation_space = self.observation_space((2*dim)**3)

        self.reset()

    def reset(self):

        self.theta = random.uniform(0, np.pi)
        self.phi = random.uniform(0, np.pi*2)

        self.ball_coordinates = [self.radius*np.sin(self.theta)*np.cos(self.phi), \
                                 self.radius*np.sin(self.theta)*np.sin(self.phi), \
                                 self.radius*np.cos(self.theta)]   #Start from an outside position

        self.state = self.get_state()
        return self.get_observation()

    def get_state(self, ball_coordinates=None):
        if ball_coordinates is None:
            ball_coordinates=self.ball_coordinates

        state = torch.zeros((2*self.size, 2*self.size, 2*self.size))

        ball_low_x = ball_coordinates[0] + self.size - .5
        ball_low_y = ball_coordinates[1] + self.size - .5
        ball_low_z = ball_coordinates[2] + self.size - .5

        fill_x = (int(ball_low_x + 1) - ball_low_x)
        fill_y = (int(ball_low_y + 1) - ball_low_y)
        fill_z = (int(ball_low_z + 1) - ball_low_z)

        state[int(ball_low_x), int(ball_low_y), int(ball_low_z)] = fill_x * fill_y * fill_z
        state[(int(ball_low_x) + 1), int(ball_low_y), int(ball_low_z)] = (1 - fill_x) * fill_y * fill_z
        state[int(ball_low_x), (int(ball_low_y) + 1), int(ball_low_z)] = fill_x * (1 - fill_y) * fill_z
        state[(int(ball_low_x) + 1), (int(ball_low_y) +1), int(ball_low_z)] = (1 - fill_x) * (1 - fill_y) * fill_z
        state[int(ball_low_x), int(ball_low_y), int(ball_low_z) + 1] = fill_x * fill_y * (1 - fill_z)
        state[(int(ball_low_x) + 1), int(ball_low_y), int(ball_low_z) + 1] = (1 - fill_x) * fill_y * (1 - fill_z)
        state[int(ball_low_x), (int(ball_low_y) + 1), int(ball_low_z) + 1] = fill_x * (1 - fill_y) * (1 - fill_z)
        state[(int(ball_low_x) + 1), (int(ball_low_y) +1), int(ball_low_z) + 1] = (1 - fill_x) * (1 - fill_y) * (1 - fill_z)

        self.state = state
        return state

    def get_observation(self):
        return self.state.flatten()

    def step(self, action, angle):

        ball_coordinates = self.ball_coordinates

        if action == 0:
            new_ball_coordinates = [ball_coordinates[0]*np.cos(angle) + ball_coordinates[1]*np.sin(angle), \
                                    - ball_coordinates[0]*np.sin(angle) + ball_coordinates[1]*np.cos(angle), \
                                    ball_coordinates[2]]
        elif action == 1:
            new_ball_coordinates = [ball_coordinates[0]*np.cos(angle) + ball_coordinates[2]*np.sin(angle), \
                                    ball_coordinates[1], \
                                    - ball_coordinates[0]*np.sin(angle) + ball_coordinates[2]*np.cos(angle)]
        elif action == 2:
            new_ball_coordinates = [ball_coordinates[0], \
                                    ball_coordinates[1]*np.cos(angle) + ball_coordinates[2]*np.sin(angle), \
                                    - ball_coordinates[1]*np.sin(angle) + ball_coordinates[2]*np.cos(angle)]
        else:
            raise Exception("INVALID ACTION")

        self.ball_coordinates = new_ball_coordinates
        self.state = self.get_state()
        return self.get_observation()
```

```python
def plot_state(state, ax):
    ax[0].pcolormesh(env.state.sum(axis=0), edgecolors='gray', linewidth=2, )
    ax[0].set_aspect('equal')
    ax[0].set_xticks([])
    ax[0].set_yticks([])

    ax[1].pcolormesh(env.state.sum(axis=1), edgecolors='gray', linewidth=2, )
    ax[1].set_aspect('equal')
    ax[1].set_xticks([])
    ax[1].set_yticks([])

    ax[2].pcolormesh(env.state.sum(axis=2), edgecolors='gray', linewidth=2, )
    ax[2].set_aspect('equal')
    ax[2].set_xticks([])
    ax[2].set_yticks([])

    return ax

n_steps = 8

env = SphereWorld()

fig, axs = plt.subplots(3, n_steps+1,figsize=(15, 5))

plot_state(env.state,axs[:,0])

for i in range(n_steps):
    action = 2
    env.step(action, np.pi/4)
    plot_state(env.state,axs[:,i+1])
```

### Latent space

**Encoder/Decoder**

Now we want to learn to represent this environment in some latent space (which we, for now, simply assume to be 3-dimensional).  We will require both an encoder and decoder.

```python
class Coder(nn.Module):

    def __init__(self, n_in, n_out, n_hid=[16], output_activation=nn.Sigmoid):
        super().__init__()

        if type(n_hid) != list:
            n_hid = [n_hid]
        n_layers = [n_in] + n_hid + [n_out]

        self.layers = []
        for i_layer, (n1, n2) in enumerate(zip(n_layers, n_layers[1:])):
            mods = [nn.Linear(n1, n2, bias=True)]
            act_fn = nn.ReLU if i_layer < len(n_layers) - 2 else output_activation
            if act_fn is not None:
                mods.append(act_fn())
            layer = nn.Sequential(*mods)
            self.layers.append(layer)

        self.layers = nn.ModuleList(self.layers)

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

class nnNorm(nn.Module):

    def __init__(self, dim=-1):
        super().__init__()
        self.dim=dim

    def forward(self, x):
        return F.normalize(x, dim=self.dim)

class Decoder(Coder):

    def __init__(self, n_in, n_out, n_hid=[32]):
        super().__init__(n_in, n_out, n_hid, output_activation=nn.Sigmoid)

class Encoder(Coder):

    def __init__(self, n_in, n_out, n_hid=[32]):
        super().__init__(n_in, n_out, n_hid, output_activation=nnNorm)

class Angles(Coder):

    def __init__(self, n_in, n_out, n_hid=[32]):
        super().__init__(n_in, n_out, n_hid, output_activation=None)

```

```python
encoder = Encoder(n_in=5*5*5, n_out=6, n_hid=[16])
print(encoder)

decoder = Decoder(n_in=6, n_out=5*5*5, n_hid=[16])
print(decoder)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

count_parameters(encoder)
```

    Encoder(
      (layers): ModuleList(
        (0): Sequential(
          (0): Linear(in_features=125, out_features=16, bias=True)
          (1): ReLU()
        )
        (1): Sequential(
          (0): Linear(in_features=16, out_features=6, bias=True)
          (1): nnNorm()
        )
      )
    )
    Decoder(
      (layers): ModuleList(
        (0): Sequential(
          (0): Linear(in_features=6, out_features=16, bias=True)
          (1): ReLU()
        )
        (1): Sequential(
          (0): Linear(in_features=16, out_features=125, bias=True)
          (1): Sigmoid()
        )
      )
    )

    2118

**Check tensor dimensions**

```python
x = torch.rand(125)
z = encoder(x)
x_recon = decoder(z)
print("Observable rep (shape = {}) : \n {}".format(x.shape, x))
print("Latent rep (shape = {}, norm = {}) : \n {}".format(z.shape, torch.norm(z), z))
print("Reconstructed Observable rep (shape = {}) : \n {}".format(x_recon.shape, x_recon))
```

    Observable rep (shape = torch.Size([125])) :
     tensor([0.7465, 0.3080, 0.6728, 0.6707, 0.2601, 0.5539, 0.2608, 0.1451, 0.9772,
            0.1851, 0.9757, 0.7232, 0.9828, 0.0409, 0.1499, 0.5088, 0.7317, 0.2052,
            0.4871, 0.9581, 0.8763, 0.5830, 0.9491, 0.2963, 0.6474, 0.6051, 0.6445,
            0.0425, 0.8856, 0.6241, 0.6790, 0.6623, 0.1144, 0.0890, 0.2951, 0.4601,
            0.1984, 0.2154, 0.2393, 0.5368, 0.8392, 0.7534, 0.8539, 0.3853, 0.6711,
            0.1581, 0.0553, 0.7864, 0.0143, 0.0531, 0.6825, 0.0017, 0.0765, 0.6904,
            0.0106, 0.5519, 0.7424, 0.5537, 0.1419, 0.5953, 0.2058, 0.0273, 0.7198,
            0.5424, 0.4027, 0.3631, 0.9479, 0.7090, 0.4094, 0.3788, 0.8239, 0.1186,
            0.8091, 0.8597, 0.0639, 0.0113, 0.1291, 0.5800, 0.2463, 0.4774, 0.8501,
            0.3880, 0.2150, 0.3972, 0.0666, 0.5490, 0.9660, 0.9858, 0.0052, 0.5190,
            0.4252, 0.4801, 0.4965, 0.1412, 0.1996, 0.3487, 0.7313, 0.3509, 0.1360,
            0.7507, 0.7039, 0.0814, 0.5088, 0.9698, 0.1575, 0.3549, 0.5117, 0.6822,
            0.4147, 0.1643, 0.4616, 0.9411, 0.1005, 0.3785, 0.5285, 0.6211, 0.7032,
            0.8068, 0.1490, 0.1894, 0.1811, 0.5365, 0.7551, 0.0856, 0.7671])
    Latent rep (shape = torch.Size([6]), norm = 0.9999999403953552) :
     tensor([ 0.0734,  0.3882,  0.0654, -0.4400, -0.3374, -0.7295],
           grad_fn=<DivBackward0>)
    Reconstructed Observable rep (shape = torch.Size([125])) :
     tensor([0.4435, 0.4840, 0.4763, 0.5231, 0.4479, 0.5144, 0.4373, 0.5287, 0.5393,
            0.5086, 0.5178, 0.5482, 0.4708, 0.4961, 0.5405, 0.4658, 0.4457, 0.4754,
            0.5614, 0.4814, 0.5040, 0.5211, 0.4365, 0.4638, 0.4744, 0.4259, 0.5092,
            0.4194, 0.5460, 0.5276, 0.5326, 0.5016, 0.5576, 0.5484, 0.4587, 0.4387,
            0.5319, 0.5185, 0.4683, 0.4921, 0.4716, 0.5211, 0.5920, 0.5665, 0.5360,
            0.5819, 0.5555, 0.5510, 0.5251, 0.5512, 0.5046, 0.5236, 0.5685, 0.5275,
            0.4731, 0.5691, 0.4600, 0.5345, 0.4504, 0.5722, 0.4617, 0.4488, 0.5266,
            0.4459, 0.5089, 0.4364, 0.4955, 0.4919, 0.4816, 0.5047, 0.5640, 0.4780,
            0.4568, 0.5409, 0.4833, 0.5416, 0.5385, 0.5882, 0.5150, 0.5079, 0.5347,
            0.5542, 0.4724, 0.5224, 0.4303, 0.4934, 0.5130, 0.4511, 0.4517, 0.5035,
            0.4977, 0.4425, 0.4801, 0.5135, 0.4641, 0.5754, 0.5213, 0.5004, 0.5510,
            0.5527, 0.5116, 0.4521, 0.4954, 0.4458, 0.5424, 0.5050, 0.5206, 0.4114,
            0.5470, 0.5452, 0.5567, 0.5531, 0.4884, 0.4721, 0.4747, 0.5567, 0.5452,
            0.5418, 0.4831, 0.5352, 0.4489, 0.5428, 0.4359, 0.4947, 0.4968],
           grad_fn=<SigmoidBackward>)

**Representation**

The crux of the matter is learning to 'represent' actions in the observation space with actions in latent space.  Here, we will do this by assuming every action is a generalized rotation in latent space, which we denote with a series of 2-dimensional rotations.

A 2-d rotation is given by:

\begin{pmatrix}
\cos(\theta) & \sin(\theta) \\
-\sin(\theta) & \cos(\theta)
\end{pmatrix}

and we denote a rotation in dimensions $i$ and $j$ of a higher dimensional space as $R_{i,j}(\theta)$.  For $i=1$, $j=4$, in a 4-dimensional space:

\begin{equation}
R_{1,4}(\theta) =
\begin{pmatrix}
\cos(\theta) & 0 & 0 & \sin(\theta) \\
0 & 1 & 0 & 0 \\
0 & 0 & 1 & 0 \\
-\sin(\theta) & 0 & 0 & \cos(\theta)
\end{pmatrix}
\end{equation}

An arbitrary rotation, denoted $g$ as I am subtly moving towards this being a group action, can then be written as:

\begin{equation}
    g(\theta_{1,2},\theta_{1,3},\dots,\theta_{n-1,n}) = \prod_{i=1}^{n-1} \prod_{j=1+1}^{n} R_{i,j}(\theta_{i,j})
\end{equation}

which has $n(n-1)/2$ free parameters (i.e. $\theta_{i,j}$'s).

To learn continuous Lie groups instead of discrete rotations, each of these $\theta_{i,j}$ is now parameterized by a neural network.

```python
class Representation():

    def __init__(self, dim=3):
        self.dim = dim
        self.params = dim*(dim-1)//2
        self.thetas = torch.autograd.Variable(torch.rand(self.params), requires_grad=True)

        self.__matrix = None

    def set_thetas(self, thetas):
        self.thetas = thetas
        self.clear_matrix()

    def clear_matrix(self):
        self.__matrix = None

    def get_matrix(self):
        if self.__matrix is None:
            k = 0
            mats = []
            for i in range(self.dim-1):
                for j in range(self.dim-1-i):
                    theta_ij = self.thetas[k]
                    k+=1
                    c, s = torch.cos(theta_ij), torch.sin(theta_ij)

                    rotation_i = torch.eye(self.dim, self.dim)
                    rotation_i[i, i] = c
                    rotation_i[i, i+j+1] = s
                    rotation_i[j+i+1, i] = -s
                    rotation_i[j+i+1, j+i+1] = c

                    mats.append(rotation_i)

            def chain_mult(l):
                if len(l)>=3:
                    return l[0]@l[1]@chain_mult(l[2:])
                elif len(l)==2:
                    return l[0]@l[1]
                else:
                    return l[0]

            self.__matrix = chain_mult(mats)

        return self.__matrix
```

**LatentWorld**

Now, for symmetry's sake, we'll also have a `LatentWorld` which acts as the environment in the latent space.

```python
class LatentWorld():

    class action_space():
        def __init__(self,n_actions):
            self.n = n_actions

        def sample(self, k=1):
            return torch.randint(0,self.n,(k,))

    class observation_space():
        def __init__(self,n_features):
            self.shape = [n_features]

    def __init__(self,
                 dim=3,
                 n_actions=3,
                 n_hid=[64]):

        self.dim = dim
        self.angle_computer = Angles(n_in = 2, n_out = dim*(dim-1)//2, n_hid=n_hid)
        self.action_space = self.action_space(n_actions)
        self.observation_space = self.observation_space(dim)

        self.representation = Representation()

    def reset(self, state):
        self.state = state
        return self.get_observation()

    def clear_representations(self):
        self.representation.clear_matrix()

    def get_representation_params(self):
        params = []
        for rep in self.action_reps:
            params.append(rep.thetas)
        return params

    def save_representations(self, path):
        if os.path.splitext(path)[-1] != '.pth':
            path += '.pth'
        return torch.save(self.angle_computer.state_dict(), path)

    def load_reprentations(self, path):
        self.representation.load_state_dict(torch.load(path))

    def get_observation(self):
        return self.state

    def step(self,action_value):
        self.representation.set_thetas(self.angle_computer(action_value))
        self.state = torch.mv(self.representation.get_matrix(), self.state)
        obs = self.get_observation()
        return obs
```

## 5. Disentanglement

***Some jargon***

We focus on learning a *disentangled* representation of the actions.

Before considering how best to do this, we want to define a metric of 'disentanglement'.  We consider the evolution of an observable (latent) vector, $x \in X$ ($z \in Z$), under the element $g \in G$ of the group of symmetries generating transformations of the object.  Then we are looking for a representation, $\rho:G \rightarrow GL(V)$, such that the transformation is linear in the latent space, i.e.
\begin{equation}
    z^{\prime} = \rho(g) \cdot z.
\end{equation}
Note, in our case, the representations are the rotation matrices we learn.

For this representation to be disentangled, it means that if there exists a subgroup decomposition of $G$
\begin{equation}
    G = G_1 \times G_2 \times \dots \times G_n,
\end{equation}
then we equivalently decompose the representation, $(\rho, G)$, into subrepresentations:
\begin{equation}
    V = V_1 \oplus V_2 \oplus \dots \oplus V_n
\end{equation}
such that the restricted subrepresentations $(\rho_{\vert G_i}, V_i)_i$ are non-trivial, and the restricted subrepresentations $(\rho_{\vert G_i}, V_j)_{j \neq i}$ are trivial.

In our context, a SphereWorld can be represented as a subgroup of $\mathrm{SO}(3)$, therefore we hope to find the disentangled representation of the three continuous rotations.

***Some practicalities***

Our intuition is that the disentangled representation acts as the identity on as many dimensions as possible.  We could attempt to enforce this with some regularization during training.  Normal weight decay won't cut it, as that tries to reduce all weights, where as what we really want to do is have all *but one* of our thetas (which corresponds to the rotation/coupling of two dimensions) to be zero.

**1. Entanglement regularisation**

So for $m$ parameters, ${\theta_1, \dots, \theta_m}$, we want to regularise with
\begin{equation}
    \sum_{i \neq j} \vert\theta_i\vert^2, \mathrm{where\ } \theta_j {=} \mathrm{max_k}({\vert\theta_k\vert}).
\end{equation}
We will also use this term as our metric of 'entanglement'.

```python
def calc_entanglement(params):
    params = params.abs().pow(2)
    return params.sum() - params.max()

params = torch.FloatTensor([1,1,0.5,0,0])
calc_entanglement(params)
```

    tensor(1.2500)

### Training

```python
obs_env = SphereWorld(dim = 5, radius = 4.)
lat_env = LatentWorld(dim = 3,
                      n_actions = obs_env.action_space.n
                     )
decoder = Decoder(n_in = sum(lat_env.observation_space.shape),
                  n_out = sum(obs_env.observation_space.shape),
                  n_hid = [64])

encoder = Encoder(n_in = sum(obs_env.observation_space.shape),
                  n_out = sum(lat_env.observation_space.shape),
                  n_hid = [64])

optimizer_dec = optim.Adam(decoder.parameters(),
                           lr=1e-2,
#                            betas=(0.9, 0.99),
                           weight_decay=0)

optimizer_enc = optim.Adam(encoder.parameters(),
                           lr=1e-2,
#                            betas=(0.9, 0.99),
                           weight_decay=0)

optimizer_rep = optim.Adam(lat_env.angle_computer.parameters(),
                           lr=1e-2,
#                            betas=(0.9, 0.99),
                           weight_decay=0)

losses = []

```

```python
n_sgd_steps = 10000
ep_steps = 5
batch_eps = 8

#Starting by learning small angles before moving on the larger angles helps
angle_max = np.pi * 2 / 5

i = 1

t_start = time.time()

temp = 0

while i < n_sgd_steps:
    if i == 4000:
        angle_max = np.pi * 4 / 5
    if i == 7000:
        angle_max = np.pi

    loss = torch.zeros(1)

    for _ in range(batch_eps):
        t_ep = -1
        while t_ep < ep_steps:
            angle = random.uniform(-angle_max, angle_max)
            if t_ep == -1:
                obs_x = obs_env.reset()
                obs_z = lat_env.reset(encoder(obs_x))
            else:
                action = random.randrange(3)
                obs_x = obs_env.step(action, angle)
                obs_z = lat_env.step(torch.tensor([action, angle]))

            t_ep += 1

            obs_x_recon = decoder(obs_z)
            (obs_x_recon==obs_x_recon.max()).float()

            loss += F.binary_cross_entropy(obs_x_recon, obs_x)
            loss += calc_entanglement(lat_env.angle_computer(torch.tensor([action, angle])))*1e-3

    loss /= (batch_eps * ep_steps)
    loss_raw = loss

    losses.append(loss_raw.item())

    optimizer_dec.zero_grad()
    optimizer_rep.zero_grad()
    optimizer_enc.zero_grad()
    loss.backward()
    optimizer_dec.step()
    optimizer_rep.step()
    optimizer_enc.step()

    # Rember to clear the cached action representations after we update the parameters!
    lat_env.clear_representations()

    i+=1

    if i%10==0:
        print("iter {} : loss={:.5f} : last 10 iters in {:.3f}s".format(
            i, loss.item(), time.time() - t_start
            ), end="\r" if i%100 else "\n")
        t_start = time.time()
```

    iter 100 : loss=0.00873 : last 10 iters in 0.956s
    iter 200 : loss=0.00840 : last 10 iters in 0.998s
    iter 300 : loss=0.00837 : last 10 iters in 1.157s
    iter 400 : loss=0.00827 : last 10 iters in 1.201s
    iter 500 : loss=0.00833 : last 10 iters in 0.954s
    iter 600 : loss=0.00839 : last 10 iters in 0.978s
    iter 700 : loss=0.00818 : last 10 iters in 0.881s
    iter 800 : loss=0.00736 : last 10 iters in 0.885s
    iter 900 : loss=0.00769 : last 10 iters in 1.069s
    iter 1000 : loss=0.00776 : last 10 iters in 0.882s
    iter 1100 : loss=0.00755 : last 10 iters in 0.900s
    iter 1200 : loss=0.00808 : last 10 iters in 0.860s
    iter 1300 : loss=0.00797 : last 10 iters in 0.859s
    iter 1400 : loss=0.00731 : last 10 iters in 0.873s
    iter 1500 : loss=0.00709 : last 10 iters in 0.995s
    iter 1600 : loss=0.00729 : last 10 iters in 1.008s
    iter 1700 : loss=0.00607 : last 10 iters in 1.166s
    iter 1800 : loss=0.00549 : last 10 iters in 0.997s
    iter 1900 : loss=0.00500 : last 10 iters in 0.996s
    iter 2000 : loss=0.00457 : last 10 iters in 1.237s
    iter 2100 : loss=0.00448 : last 10 iters in 0.981s
    iter 2200 : loss=0.00482 : last 10 iters in 0.958s
    iter 2300 : loss=0.00462 : last 10 iters in 1.115s
    iter 2400 : loss=0.00416 : last 10 iters in 0.855s
    iter 2500 : loss=0.00447 : last 10 iters in 0.899s
    iter 2600 : loss=0.00377 : last 10 iters in 0.949s
    iter 2700 : loss=0.00400 : last 10 iters in 0.944s
    iter 2800 : loss=0.00375 : last 10 iters in 0.863s
    iter 2900 : loss=0.00375 : last 10 iters in 0.891s
    iter 3000 : loss=0.00375 : last 10 iters in 0.864s
    iter 3100 : loss=0.00364 : last 10 iters in 0.961s
    iter 3200 : loss=0.00397 : last 10 iters in 1.603s
    iter 3300 : loss=0.00366 : last 10 iters in 1.252s
    iter 3400 : loss=0.00369 : last 10 iters in 0.966s
    iter 3500 : loss=0.00353 : last 10 iters in 0.887s
    iter 3600 : loss=0.00369 : last 10 iters in 0.938s
    iter 3700 : loss=0.00356 : last 10 iters in 0.918s
    iter 3800 : loss=0.00372 : last 10 iters in 0.868s
    iter 3900 : loss=0.00352 : last 10 iters in 0.868s
    iter 4000 : loss=0.00344 : last 10 iters in 1.007s
    iter 4100 : loss=0.00434 : last 10 iters in 1.010s
    iter 4200 : loss=0.00370 : last 10 iters in 0.914s
    iter 4300 : loss=0.00371 : last 10 iters in 0.912s
    iter 4400 : loss=0.00355 : last 10 iters in 0.880s
    iter 4500 : loss=0.00350 : last 10 iters in 0.930s
    iter 4600 : loss=0.00352 : last 10 iters in 0.868s
    iter 4700 : loss=0.00401 : last 10 iters in 0.916s
    iter 4800 : loss=0.00346 : last 10 iters in 0.881s
    iter 4900 : loss=0.00350 : last 10 iters in 0.868s
    iter 5000 : loss=0.00345 : last 10 iters in 0.845s
    iter 5100 : loss=0.00349 : last 10 iters in 0.894s
    iter 5200 : loss=0.00337 : last 10 iters in 0.864s
    iter 5300 : loss=0.00333 : last 10 iters in 0.907s
    iter 5400 : loss=0.00343 : last 10 iters in 0.931s
    iter 5500 : loss=0.00581 : last 10 iters in 0.960s
    iter 5600 : loss=0.00356 : last 10 iters in 1.423s
    iter 5700 : loss=0.00345 : last 10 iters in 1.876s
    iter 5800 : loss=0.00334 : last 10 iters in 1.434s
    iter 5900 : loss=0.00356 : last 10 iters in 1.442s
    iter 6000 : loss=0.00343 : last 10 iters in 1.393s
    iter 6100 : loss=0.00344 : last 10 iters in 0.848s
    iter 6200 : loss=0.00368 : last 10 iters in 0.863s
    iter 6300 : loss=0.00334 : last 10 iters in 1.266s
    iter 6400 : loss=0.00340 : last 10 iters in 0.966s
    iter 6500 : loss=0.00320 : last 10 iters in 1.042s
    iter 6600 : loss=0.00342 : last 10 iters in 1.074s
    iter 6700 : loss=0.00361 : last 10 iters in 1.237s
    iter 6800 : loss=0.00374 : last 10 iters in 0.877s
    iter 6900 : loss=0.00328 : last 10 iters in 1.039s
    iter 7000 : loss=0.00313 : last 10 iters in 0.882s
    iter 7100 : loss=0.00444 : last 10 iters in 0.994s
    iter 7200 : loss=0.00338 : last 10 iters in 1.122s
    iter 7300 : loss=0.00337 : last 10 iters in 0.939s
    iter 7400 : loss=0.00380 : last 10 iters in 1.519s
    iter 7500 : loss=0.00337 : last 10 iters in 1.227s
    iter 7600 : loss=0.00325 : last 10 iters in 1.615s
    iter 7700 : loss=0.00342 : last 10 iters in 1.421s
    iter 7800 : loss=0.00353 : last 10 iters in 0.886s
    iter 7900 : loss=0.00348 : last 10 iters in 1.232s
    iter 8000 : loss=0.00399 : last 10 iters in 1.118s
    iter 8100 : loss=0.00344 : last 10 iters in 0.859s
    iter 8200 : loss=0.00342 : last 10 iters in 0.855s
    iter 8300 : loss=0.00327 : last 10 iters in 0.878s
    iter 8400 : loss=0.00315 : last 10 iters in 0.855s
    iter 8500 : loss=0.00379 : last 10 iters in 1.099s
    iter 8600 : loss=0.00339 : last 10 iters in 1.005s
    iter 8700 : loss=0.00348 : last 10 iters in 1.065s
    iter 8800 : loss=0.00339 : last 10 iters in 1.066s
    iter 8900 : loss=0.00349 : last 10 iters in 1.237s
    iter 9000 : loss=0.00317 : last 10 iters in 0.916s
    iter 9100 : loss=0.00347 : last 10 iters in 1.120s
    iter 9200 : loss=0.00315 : last 10 iters in 0.857s
    iter 9300 : loss=0.00399 : last 10 iters in 0.868s
    iter 9400 : loss=0.00318 : last 10 iters in 1.083s
    iter 9500 : loss=0.00317 : last 10 iters in 0.850s
    iter 9600 : loss=0.00339 : last 10 iters in 1.345s
    iter 9700 : loss=0.00347 : last 10 iters in 1.166s
    iter 9800 : loss=0.00363 : last 10 iters in 1.368s
    iter 9900 : loss=0.00320 : last 10 iters in 0.930s
    iter 10000 : loss=0.00362 : last 10 iters in 1.116s

```python
exp_name = "ent_reg_1e-2_orthNorm_reg_1e-2_20steps_16batch"
# exp_name = "testing"

save_folder = os.path.join("_data/decoder_only", exp_name)
train_folder = os.path.join(save_folder,"train")
test_folder = os.path.join(save_folder,"test")

mk_dir(save_folder)
mk_dir(train_folder)
mk_dir(test_folder)
```

    dir already exists:  _data/decoder_only/ent_reg_1e-2_orthNorm_reg_1e-2_20steps_16batch
    dir already exists:  _data/decoder_only/ent_reg_1e-2_orthNorm_reg_1e-2_20steps_16batch/train
    dir already exists:  _data/decoder_only/ent_reg_1e-2_orthNorm_reg_1e-2_20steps_16batch/test

**Plot training curve**

```python
with plt.style.context('seaborn-paper', after_reset=True):

    fig, (ax1) = plt.subplots(1, 1, figsize=(8, 8), gridspec_kw={"wspace":0.3})

    window = 25
    avg_mask = np.ones(window) / window

    ax1.plot(np.convolve(range(len(losses)), avg_mask, 'valid'),
             np.convolve(np.log(losses), avg_mask, 'valid'),
             linewidth=0.75,
             alpha=0.8)
#     ax1.set_yscale("log")

    ax1.set_xlabel("Num. parameter updates")
    ax1.set_ylabel("Log reconstruction loss")

    fig_fname = os.path.join(train_folder, "training_curves")
    plt.savefig(fig_fname + ".pdf", bbox_inches='tight')
    plt.savefig(fig_fname + ".png", bbox_inches='tight')
```

```python
torch.save(decoder.state_dict(), os.path.join(save_folder,"decoder.pth"))
torch.save(optimizer_dec.state_dict(), os.path.join(train_folder,"optimizer_dec.pth"))
torch.save(optimizer_rep.state_dict(), os.path.join(train_folder,"optimizer_rep.pth"))
```

### Testing

**1) Representations learned for specific values of the rotations**

```python
width=0.75

plt_lim = 0.22
titles = ["-0.2", "-0.1", "0", "0.1", "0.2"]

with plt.style.context('seaborn-paper', after_reset=True):

    fig, axs = plt.subplots(3, 5, figsize=(15, 5), gridspec_kw={"wspace":0.4})

    for i in range(15):
        thetas = lat_env.angle_computer(torch.tensor([i//5, (-0.2+(i%5)/10.)*2*np.pi])).detach()
        print(thetas)
        x = np.arange(len(thetas))
        axs[i//5,i%5].bar(x - width/2, thetas/(2*np.pi), width, label='Rep {}'.format(i))

        axs[i//5,i%5].set_yticks([-0.2, 0., 0.2])
        axs[i//5,i%5].set_xticks(x-0.25)
        axs[i//5,i%5].set_xticklabels(["12","13","23"])
        axs[i//5,i%5].set_xlabel("$ij$")

        axs[i//5,i%5].set_ylim(-plt_lim,plt_lim)

        axs[i//5,i%5].set_title(titles[i%5])

    axs[0,0].set_ylabel(r"$x : \theta / 2\pi$" )
    axs[1,0].set_ylabel(r"$y : \theta / 2\pi$" )
    axs[2,0].set_ylabel(r"$y : \theta / 2\pi$" )

    fig_fname = os.path.join(test_folder, "thetas")

    plt.savefig(fig_fname + ".pdf", bbox_inches='tight')
    plt.savefig(fig_fname + ".png", bbox_inches='tight')
```

**2) Learned representations as a function of rotation angle**

```python
fig, axs = plt.subplots(1, 3, figsize=(18, 5), gridspec_kw={"wspace":0.4})

x = np.linspace(-np.pi, np.pi, 500)

titles = ["rotation around $x$", "rotation around $y$", "rotation around $z$"]

for i in range(3):
    y0 = []
    y1 = []
    y2 = []
    for k in range(500):
        y = lat_env.angle_computer(torch.tensor([i, x[k]])).detach().numpy()
        y0 += [y[0]]
        y1 += [y[1]]
        y2 += [y[2]]
    a = axs[i].plot(x,y0,c="r",linewidth=2)[0]
    b = axs[i].plot(x,y1,c="b",linewidth=2)[0]
    c = axs[i].plot(x,y2,c="g",linewidth=2)[0]
    axs[i].set_xlabel("angle of the rotation", fontsize=15)
    axs[i].set_title(titles[i], fontsize=15)
    axs[i].tick_params(labelsize=12)
axs[0].set_ylabel(r"$\theta_{ij}$", fontsize=20)
fig.legend((a,b,c),(r'ij=12',r'ij=13',r'ij=23'),loc="center right", fontsize=18)
fig.show()

```

**3) See predictions made by network**

```python
def plot_state(obs, ax, one_hot=True):
    if one_hot:
        obs=(obs == obs.max())
    ax.pcolormesh(obs.reshape(2*obs_env.size, 2*obs_env.size, 2*obs_env.size).sum(axis=2), edgecolors='gray', linewidth=2)
    ax.set_aspect('equal')
    ax.set_xticks([])
    ax.set_yticks([])

    return ax

n_steps = 10

fig, (ax1,ax2) = plt.subplots(2, 1)

ax1.set_title("Ground truth")
ax2.set_title("Reconstruction")

for i in range(n_steps+1):

    if i==0:
        action = "N\A"
        obs_x = obs_env.reset()
        obs_z = lat_env.reset(encoder(obs_x))
    else:
        action = random.randrange(3)
        angle = random.uniform(-angle_max, angle_max)
        obs_x = obs_env.step(action, angle)
        obs_z = lat_env.step(torch.tensor([action, angle]))

    obs_x_recon = decoder(obs_z)

    fig.suptitle('step {} : last action = {} in the x-y plane'.format(i, action), fontsize=16)

    plot_state(obs_x.detach().numpy(),ax1)
    plot_state(obs_x_recon.detach().numpy(),ax2, one_hot=False)

    fig_fname = os.path.join(test_folder, "step_{}".format(i))

    plt.savefig(fig_fname + ".pdf", bbox_inches='tight')
    plt.savefig(fig_fname + ".png", bbox_inches='tight')

    display.clear_output(wait=True)
    display.display(plt.gcf())
    time.sleep(1.)

display.clear_output(wait=False)
```

```python

```


# DL Neuromatch Academy: Week 1, Day 2, Tutorial 3
# Deep Linear Neural Networks

__Content creators:__ Andrew Saxe, Saeed Salehi, Vladimir Haltakov

__Content reviewers:__ Polina Turishcheva, Atnafu Lambebo, Yu-Fang Yang

__Content editors:__ Anoop Kulkarni

__Production editors:__ Khalid Almubarak, , Spiros Chavlis

---
#Tutorial Objectives

* Deep linear neural networks
* Learning dynamics and singular value decomposition
* Representational Similarity Analysis
* Illusory correlations & ethics.

```python
#@markdown Tutorial slides
# you should link the slides for all tutorial videos here (we will store pdfs on osf)

from IPython.display import HTML
HTML('')
```

---
# Setup

```python
# Imports

import numpy as np
import matplotlib.pyplot as plt
import random
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.optim as optim

from tqdm.notebook import tqdm, trange
import time
from math import sqrt

! pip install treelib --quiet
from treelib import Node, Tree
```

```python
# @title Figure settings
# import ipywidgets as widgets
# from ipywidgets import interact, fixed, HBox, Layout, VBox, interactive, Label
# from ipywidgets import interact, IntSlider, FloatSlider, interact_manual
# from mpl_toolkits.mplot3d import Axes3D
from matplotlib import gridspec
from ipywidgets import interact, IntSlider, FloatSlider, interact_manual, fixed
from ipywidgets import FloatLogSlider, HBox, Layout, VBox, interactive, Label
from ipywidgets import interactive_output

import warnings
warnings.filterwarnings("ignore")

%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
#@title Plotting functions

def plot_x_y_hier_data(im1, im2, subplot_ratio=[1, 2]):
  fig = plt.figure(figsize=(12, 5))
  gs = gridspec.GridSpec(1, 2, width_ratios=subplot_ratio)
  ax0 = plt.subplot(gs[0])
  ax1 = plt.subplot(gs[1])
  ax0.imshow(im1, cmap="cool")
  ax1.imshow(im2, cmap="cool")
  # plt.suptitle("The whole dataset as imshow plot", y=1.02)
  ax0.set_title("Labels of all samples")
  ax1.set_title("Features of all samples")
  ax0.set_axis_off()
  ax1.set_axis_off()
  plt.tight_layout()
  plt.show()

def plot_x_y_hier_one(im1, im2, subplot_ratio=[1, 2]):
  fig = plt.figure(figsize=(12, 1))
  gs = gridspec.GridSpec(1, 2, width_ratios=subplot_ratio)
  ax0 = plt.subplot(gs[0])
  ax1 = plt.subplot(gs[1])
  ax0.imshow(im1, cmap="cool")
  ax1.imshow(im2, cmap="cool")
  ax0.set_title("Labels of a single sample")
  ax1.set_title("Features of a single sample")
  ax0.set_axis_off()
  ax1.set_axis_off()
  plt.tight_layout()
  plt.show()

def plot_tree_data(im1, im2, label_list):
  im1_dim1, im1_dim2 = im1.shape

  fig = plt.figure(figsize=(12, 5))
  gs = gridspec.GridSpec(1, 2, width_ratios=[1, 2])
  ax0 = plt.subplot(gs[0])
  ax1 = plt.subplot(gs[1])
  ax0.imshow(im1, cmap="cool")
  ax1.imshow(im2[:, -im1_dim1*2:], cmap="cool", vmin=0.0, vmax=1.0)
  ax0.set_title("all the Labels")
  ax1.set_title("last {} Features".format(-im1_dim1*2))
  ax0.set_yticks(ticks=np.arange(im1_dim1))
  ax0.set_yticklabels(labels=label_list)
  ax0.set_xticks(ticks=np.arange(im1_dim1))
  ax0.set_xticklabels(labels=item_names, rotation='vertical')
  ax1.set_axis_off()
  plt.tight_layout()
  plt.show()

def plot_loss(loss_array, title="Training loss (Mean Squared Error)", c="r"):
  plt.figure(figsize=(9, 5))
  plt.plot(loss_array, color=c)
  plt.xlabel("Epoch")
  plt.ylabel("MSE")
  plt.title(title)
  plt.show()

def plot_loss_sv(loss_array, sv_array):
  n_sing_values = sv_array.shape[1]
  sv_array = sv_array / np.max(sv_array)
  cmap = plt.cm.get_cmap("Set1", n_sing_values)

  _, (plot1, plot2) = plt.subplots(2, 1, sharex=True, figsize=(10, 10))
  plot1.set_title("Training loss (Mean Squared Error)")
  plot1.plot(loss_array, color='r')

  plot2.set_title("Evolution of singular values (modes)")
  for i in range(n_sing_values):
    plot2.plot(sv_array[:, i], c=cmap(i))
  plot2.set_xlabel("Epoch")
  plt.show()

def plot_loss_sv_twin(loss_array, sv_array):
  n_sing_values = sv_array.shape[1]
  sv_array = sv_array / np.max(sv_array)
  cmap = plt.cm.get_cmap("winter", n_sing_values)

  fig = plt.figure(figsize=(11, 6))
  ax1 = plt.gca()
  ax1.set_title("Learning Dynamics")
  ax1.set_xlabel("Epoch")
  ax1.set_ylabel("Mean Squared Error", c='r')
  ax1.tick_params(axis='y', labelcolor='r')
  ax1.plot(loss_array, color='r')

  ax2 = ax1.twinx()
  ax2.set_ylabel("Singular values (modes)", c='b')
  ax2.tick_params(axis='y', labelcolor='b')
  for i in range(n_sing_values):
    ax2.plot(sv_array[:, i], c=cmap(i))

  fig.tight_layout()
  plt.show()

def plot_ills_sv_twin(ill_array, sv_array):
  n_sing_values = sv_array.shape[1]
  sv_array = sv_array / np.max(sv_array)
  cmap = plt.cm.get_cmap("winter", n_sing_values)

  fig = plt.figure(figsize=(11, 6))
  ax1 = plt.gca()
  ax1.set_title("Network evolution in learning the Illusory Correlations")
  ax1.set_xlabel("Epoch")
  ax1.set_ylabel("Illusory Correlations", c='r')
  ax1.tick_params(axis='y', labelcolor='r')
  ax1.plot(ill_array, color='r', linewidth=3)
  ax1.set_ylim(-0.05, 1.0)

  ax2 = ax1.twinx()
  ax2.set_ylabel("Singular values (modes)", c='b')
  ax2.tick_params(axis='y', labelcolor='b')
  for i in range(n_sing_values):
    ax2.plot(sv_array[:, i], c=cmap(i))

  fig.tight_layout()
  plt.show()

def plot_loss_sv_rsm(loss_array, sv_array, rsm_array, i_ep):
  rsm_array = rsm_array / np.max(rsm_array, axis=0)
  sv_array = sv_array / np.max(sv_array)

  n_sing_values = sv_array.shape[1]
  cmap = plt.cm.get_cmap("winter", n_sing_values)

  fig = plt.figure(figsize=(15, 5))
  gs = gridspec.GridSpec(1, 2, width_ratios=[2, 1])

  ax0 = plt.subplot(gs[1])
  ax0.yaxis.tick_right()
  ax0.imshow(rsm_array[i_ep], cmap="Purples", vmin=0.0, vmax=1.1)
  ax0.set_title("RSM at epoch {}".format(i_ep), fontsize=16)
  # ax0.set_axis_off()
  ax0.set_yticks(ticks=np.arange(n_sing_values))
  ax0.set_yticklabels(labels=item_names)
  # ax0.set_xticks([])
  ax0.set_xticks(ticks=np.arange(n_sing_values))
  ax0.set_xticklabels(labels=item_names, rotation='vertical')

  ax1 = plt.subplot(gs[0])
  ax1.set_title("Learning Dynamics", fontsize=16)
  ax1.set_xlabel("Epoch")
  ax1.set_ylabel("Mean Squared Error", c='r')
  ax1.tick_params(axis='y', labelcolor='r')
  ax1.plot(loss_array, color='r')
  ax1.axvspan(i_ep-2, i_ep+2, alpha=0.2, color='m')

  ax2 = ax1.twinx()
  ax2.set_ylabel("Singular values", c='b')
  ax2.tick_params(axis='y', labelcolor='b')
  for i in range(n_sing_values):
    ax2.plot(sv_array[:, i], c=cmap(i))

  plt.show()

class SimpleTree:
  def __init__(self, plot=False):
    tree = Tree()
    tree.create_node("Living things", 0)
    tree.create_node("Animal", 1, parent=0)
    tree.create_node("Plant", 2, parent=0)
    tree.create_node("Fish", 3, parent=1)
    tree.create_node("Bird", 4, parent=1)
    tree.create_node("Flower", 5, parent=2)
    tree.create_node("Tree", 6, parent=2)
    tree.create_node("Goldfish", 7, parent=3)
    tree.create_node("Tuna", 8, parent=3)
    tree.create_node("Robin", 9, parent=4)
    tree.create_node("Canary", 10, parent=4)
    tree.create_node("Rose", 11, parent=5)
    tree.create_node("Daisy", 12, parent=5)
    tree.create_node("Pine", 13, parent=6)
    tree.create_node("Oak", 14, parent=6)
    self.tree = tree
    if plot: self.plot()

  def plot(self):
    self.tree.show(line_type="ascii-em")

  def rename(self, old, new, plot=False):
    for nodes in self.tree.all_nodes():
      if nodes.tag == old:
        nodes.tag = new
        break
    if plot: self.plot()
```

```python
#@title Helper functions

seed = 2015  # LeCun, Bengio, & Hinton (2015)
random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

class VariableDepthWidth(nn.Module):
  def __init__(self, in_dim, out_dim, hid_dims=[], gamma=1e-12):
    """Variable depth linear network

    Args:
      in_dim (int): input dimension
      out_dim (int): ouput dimension
      hid_dims (list): a list, containing the number of neurons in each hidden layer
        default is empty list (`[]`) for linear regression.
        example: For 2 hidden layers, first with 5 and second with 7 neurons,
                 we use: `hid_dims = [5, 7]`
    """
    super().__init__()
    assert isinstance(in_dim, int)
    assert isinstance(out_dim, int)
    assert isinstance(hid_dims, list)
    n_hidden_layers = len(hid_dims)  # number of hidden layers
    layers = OrderedDict()

    if n_hidden_layers == 0:  # linear regression
      layers["map"] = nn.Linear(in_dim, out_dim, bias=False)

    else:  # shallow and deep linear neural net
      layers["in->"] = nn.Linear(in_dim, hid_dims[0], bias=False)

      for i in range(n_hidden_layers-1):  # creating hidden layers
        layers["hid {}".format(i+1)] = nn.Linear(hid_dims[i],
                                                    hid_dims[i+1],
                                                    bias=False)

      layers["->out"] = nn.Linear(hid_dims[-1], out_dim, bias=False)

    for k in layers:  # re-initialization of the weights
      sigma = gamma / sqrt(layers[k].weight.shape[0] + layers[k].weight.shape[1])
      nn.init.normal_(layers[k].weight, std=sigma)

    self.layers = nn.Sequential(layers)

  def forward(self, input_tensor):
    """Forward pass
    """
    return self.layers(input_tensor)

def build_tree(n_levels, n_branches, probability, to_np_array=True):
  """Builds a tree
  """
  assert 0.0 <= probability <= 1.0

  tree = {}

  tree["level"] = [0]
  for i in range(1, n_levels+1):
    tree["level"].extend([i]*(n_branches**i))

  tree["pflip"] = [probability]*len(tree["level"])

  tree["parent"] = [None]
  k = len(tree["level"])-1
  for j in range(k//n_branches):
    tree["parent"].extend([j]*n_branches)

  if to_np_array:
    tree["level"] = np.array(tree["level"])
    tree["pflip"] = np.array(tree["pflip"])
    tree["parent"] = np.array(tree["parent"])

  return tree

def sample_from_tree(tree, n):
  """ Generates n samples from a tree
  """
  items = [i for i, v in enumerate(tree["level"]) if v == max(tree["level"])]
  n_items = len(items)
  x = np.zeros(shape=(n, n_items))
  rand_temp = np.random.rand(n, len(tree["pflip"]))
  flip_temp = np.repeat(tree["pflip"].reshape(1, -1), n, 0)
  samp = (rand_temp > flip_temp) * 2 - 1

  for i in range(n_items):
    j = items[i]
    prop = samp[:, j]
    while tree["parent"][j] is not None:
      j = tree["parent"][j]
      prop = prop * samp[:, j]
    x[:, i] = prop.T
  return x

def generate_hsd():
  # building the tree
  n_branches = 2  # 2 branches at each node
  probability = .15  # flipping probability
  n_levels = 3  # number of levels (depth of tree)
  tree = build_tree(n_levels, n_branches, probability, to_np_array=True)
  tree["pflip"][0] = 0.5
  n_samples = 10000 # Sample this many features

  tree_labels = np.eye(n_branches**n_levels)
  tree_features = sample_from_tree(tree, n_samples).T
  return tree_labels, tree_features

def linear_regression(X, Y):
  """Analytical Linear regression

  """
  assert isinstance(X, np.ndarray)
  assert isinstance(Y, np.ndarray)
  M, Dx = X.shape
  N, Dy = Y.shape
  assert Dy == Dy
  W = Y @ X.T @ np.linalg.inv(X @ X.T)
  return W

# #@markdown Run this cell to define the train function!

def train_svd_rsa_track(model, in_features, out_features, n_epochs, lr, ill_i=0):
  """Training function

  Args:
    model (torch nn.Module): the neural network
    in_features (torch.Tensor): features (input) with shape `torch.Size([batch_size, input_dim])`
    out_features (torch.Tensor): targets (labels) with shape `torch.Size([batch_size, output_dim])`
    n_epochs (int): number of training epochs
    lr (float): learning rate
    ill_i (int): index of illusory feature

  Returns:
    np.ndarray: record (evolution) of losses
    np.ndarray: record (evolution) of singular values
    np.ndarray: record (evolution) of representational similarity matrices
    np.ndarray: record of network prediction for the last feature
  """

  assert in_features.shape[0] == out_features.shape[0]
  optimizer = optim.SGD(model.parameters(), lr=lr)
  criterion = nn.MSELoss()
  xd = in_features.shape[1]
  loss_record = []  # losses
  sv_record = []  # singular values
  rsm_record = []  # represent sim mats
  pred_record = []  # network prediction for the last feature

  for i in range(n_epochs):
    y_pred = model(in_features)  # forward pass
    loss = criterion(y_pred, out_features)  # calculating the loss
    optimizer.zero_grad()  # reset all the graph gradients to zero
    loss.backward()  # back propagation of the error
    optimizer.step()  # gradient step

    # calculating the W_tot by multiplying all layers' weights
    W_tot = model.layers[-1].weight.detach()  # starting from the last layer
    for i in range(2, len(model.layers)+1):
      W_tot = W_tot @ model.layers[-i].weight.detach()
    U, Σ, V = torch.svd(W_tot)  # performing the SVD!

    # calculating representational similarity matrix
    H1 = model.layers[0].weight.detach() @ in_features
    RSM = H1.T @ H1

    # network prediction of ill_i in_feature for the last feature
    ill_pred = y_pred[ill_i, -1].detach().numpy()

    loss_record.append(loss.item())
    sv_record.append(Σ.numpy())
    rsm_record.append(RSM.numpy())
    pred_record.append(ill_pred)

  return np.array(loss_record), np.array(sv_record), np.array(rsm_record), np.array(pred_record)

def add_feature(existing_features, new_feature):
  assert isinstance(existing_features, np.ndarray)
  assert isinstance(new_feature, list)
  new_feature = np.array([new_feature]).T
  # return np.hstack((tree_features, new_feature*2-1))
  return np.hstack((tree_features, new_feature))
```

---
# Section 0: Prelude

## Exercise 0: Variable depth and width LNN
Throughout this tutorial, we will need several neural nets with different depth and width. So first, let's create a model with variable depth and width.

This can be easily done by using [`OrderedDict()`](https://docs.python.org/3/library/collections.html#collections.OrderedDict) and `nn.Sequential` function. The model is defined by its input and output dimensions, and a list containing the width of each hidden layer. If the list is left empty, the neural network will perform a linear regression (coming up next). We also exclude the `bias` from all the layers.

We also take over the initialization. In PyTorch, we can use [`nn.init`](https://pytorch.org/docs/stable/nn.init.html) to initialize tensors from a given distribution. Here, we sample the weights from the following distribution:

$$\mathcal{N}\left(\mu=0, ~~\sigma=\gamma \sqrt{\dfrac{1}{n_{in} + n_{out}}} \right)$$

where $\gamma$ is given as an argument. the Underscore ("_") in `nn.init.random_` and other functions, denote "in-place" operation. Note that `nn.Linear` layers are initialized at definition, so we re-initialize them.

```python
class VariableDepthWidthExercise(nn.Module):
  def __init__(self, in_dim, out_dim, hid_dims=[], gamma=1e-12):
    """Variable depth linear network

    Args:
      in_dim (int): input dimension
      out_dim (int): ouput dimension
      hid_dims (list): a list, containing the number of neurons in each hidden layer
        default is empty list (`[]`) for linear regression.
        example: For 2 hidden layers, first with 5 and second with 7 neurons,
                 we use: `hid_dims = [5, 7]`
    """
    super().__init__()
    assert isinstance(in_dim, int)
    assert isinstance(out_dim, int)
    assert isinstance(hid_dims, list)
    n_hidden_layers = len(hid_dims)  # number of hidden layers
    layers = OrderedDict()

    if n_hidden_layers == 0:  # linear regression
      layers["map"] = nn.Linear(in_dim, out_dim, bias=False)

    else:  # shallow and deep linear neural net
      layers["in->"] = nn.Linear(in_dim, hid_dims[0], bias=False)

      for i in range(n_hidden_layers-1):  # creating hidden layers
        #################################################
        ## Complete the hidden loop of DeepLNNExercise class
        # Complete the function and remove or comment the line below
        raise NotImplementedError("Network model `DeepLNNExercise`")
        #################################################

    for k in layers:  # re-initialization of the weights
      sigma = gamma / sqrt(layers[k].weight.shape[0] + layers[k].weight.shape[1])
      nn.init.normal_(layers[k].weight, std=sigma)

    self.layers = nn.Sequential(layers)

  def forward(self, input_tensor):
    """Forward pass
    """
    return self.layers(input_tensor)

# # Uncomment and run
# print("Deep LNN:\n",
#       VariableDepthWidthExercise(64, 100, [32, 16, 16, 32]))

# print("\nLinear Regression model:\n",
#       VariableDepthWidthExercise(64, 100,[]))
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content-dl/tree/main//tutorials/W1D2_LinearNN/solutions/W1D2_Tutorial3_Solution_cfc9c90c.py)

We have prepare the training function (very similar to that of tutorial 1) for you already. Just look check if everything is okay.

```python
def train(model, in_features, out_features, n_epochs, lr,
          criterion=None, optimizer=None, show_progress_bar=False):
  """Training function

  Args:
    model (torch nn.Module): the neural network
    in_features (torch.Tensor): features (input) with shape `torch.Size([batch_size, input_dim])`
    out_features (torch.Tensor): targets (labels) with shape `torch.Size([batch_size, output_dim])`
    n_epochs (int): number of training epochs
    criterion (function): loss function (default 'nn.MSELoss()')
    optimizer(function): optimizer (default 'optim.SGD')
    lr(float): learning rate

  Returns:
    list: record (evolution) of losses
  """

  assert in_features.shape[0] == out_features.shape[0]
  loss_record = []  # for recoding losses

  if optimizer is None:
      optimizer = optim.SGD(model.parameters(), lr=lr)
  if criterion is None:
      criterion = nn.MSELoss()

  model.train()  # we first put the model in training mode

  for i in range(n_epochs):
    y_pred = model(in_features)  # forward pass
    loss = criterion(y_pred, out_features)  # calculating the loss
    optimizer.zero_grad()  # reset all the graph gradients to zero
    loss.backward()  # back propagation of the error
    optimizer.step()  # gradient step

    loss_record.append(loss.item())

  model.eval()  # putting the model to evaluation mode

  return loss_record
```

---
# Section 00: Analytical Linear Regression

Linear regression is a relatively simple optimization problem. Unlike most other models that we will see in this course, linear regression for mean squared loss can be solved analytically.

For $D$ samples (batch size), $\mathbf{X} \in \mathbb{R}^{M \times D}$, and $\mathbf{Y} \in \mathbb{R}^{N \times D}$, the goal of linear regression is to find $\mathbf{W} \in \mathbb{R}^{N \times M}$ such that:

$$\mathbf{Y} = \mathbf{W} ~ \mathbf{X} $$

Given the Squared Error loss function, we have:

\begin{equation}
Loss(\mathbf{W}) = ||\mathbf{Y} - \mathbf{W} ~ \mathbf{X}||^2
\end{equation}

So, using matrix notation, the optimization problem is given by:

\begin{align}
\mathbf{W^{*}} &= \underset{\mathbf{W}}{\mathrm{argmin}} \left( Loss (\mathbf{W})  \right) \\
 &= \underset{\mathbf{W}}{\mathrm{argmin}} \left( ||\mathbf{Y} - \mathbf{W} ~ \mathbf{X}||^2  \right) \\
&= \underset{\mathbf{W}}{\mathrm{argmin}} \left( \left( \mathbf{Y} - \mathbf{W} ~ \mathbf{X}\right)^{\top} \left( \mathbf{Y} - \mathbf{W} ~ \mathbf{X}\right) \right)
\end{align}

To solve the minimization problem, we can simply set the derivative of the loss with respect to $\mathbf{W}$ to zero.

\begin{equation}
\dfrac{\partial Loss}{\partial \mathbf{W}} = 0
\end{equation}

Assuming that $\mathbf{X}\mathbf{X}^{\top}$ is full-rank, and thus it is invertible we can write:

\begin{equation}
\mathbf{W}^{\mathbf{*}} = \mathbf{Y} \mathbf{X}^{\top} \left( \mathbf{X}  \mathbf{X}^{\top} \right) ^{-1}
\end{equation}

## Exercise 00: Analytical solution to LR

Complete the function `linear_regression` for finding the analytical solution to linear regression.

```python
def linear_regression_exercise(X, Y):
  """Analytical Linear regression

  Args:
    X (np.ndarray): design matrix
    Y (np.ndarray): target ouputs

  return:
    np.ndarray: estimated weights (mapping)
  """
  assert isinstance(X, np.ndarray)
  assert isinstance(Y, np.ndarray)
  M, Dx = X.shape
  N, Dy = Y.shape
  assert Dy == Dy

  #################################################
  ## Complete the linear_regression_exercise function
  # Complete the function and remove or comment the line below
  raise NotImplementedError("Linear Regression `linear_regression_exercise`")
  #################################################
  W = ...

  return W

W_true = np.random.randint(low=0, high=10, size=(3, 3)).astype(float)

X_train = np.random.rand(3, 37)  # 37 samples
noise = np.random.normal(scale=0.01, size=(3, 37))
Y_train = W_true @ X_train + noise

# # Uncomment and run
# W_estimate = linear_regression_exercise(X_train, Y_train)

# print("True weights:\n", W_true)
# print("\nEstimated weights:\n", np.round(W_estimate, 1))
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content-dl/tree/main//tutorials/W1D2_LinearNN/solutions/W1D2_Tutorial3_Solution_ba1e16d4.py)

---
# Section 1: Deep Linear Neural Nets

```python
#@title Video 1: Representation Learning (Intro)
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="MRPy6uZRxms", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

So far depth just seems to slow down the learning. And we know that a single nonlinear hidden layer (given enough number of neurons) has the potential to approximate any function. So it's seems fair to ask: **What is depth good for**? One reason can be that, shallow nonlinear neural networks hardly meet their true potential in practice.

In the contrast, deep neural nets are often surprisingly powerful in learning complex functions without sacrificing the generalization. A core intuition behind deep learning is that deep nets derive their power through learning internal representations. How does this work? To address representation learning, we have to go beyond the 1D chain, to a deep Linear Neual Network (LNN).

For this and the next couple of exercises, we use syntactically generated hierarchically structured data through a *branching diffusion process* (see [this reference](https://www.pnas.org/content/pnas/suppl/2019/05/16/1820226116.DCSupplemental/pnas.1820226116.sapp.pdf) for more details).

 hierarchically structured data

## Exercise 1: Training a deep LNN

This is a rather simple exercise. We will generate some hierarchically structured data, instantiate a LNN from `VariableDepthWidth` class and train it on the data.

**Important note**:

* Datasets are often generated as `numpy.ndarray`, and pass to PyTorch, which needs `torch.Tensor` for training. You can use `torch.tensor(toy_data).float()` for the "conversion" of float datatype.

```python
#@markdown #### Run to generate and visualize training samples from tree

tree_labels, tree_features = generate_hsd()
item_names = ['Goldfish', 'Tuna', 'Robin', 'Canary', 'Rose', 'Daisy', 'Pine', 'Oak']
plot_tree_data(tree_labels, tree_features, item_names)

# dimensions
print()
print("---------------------------------------------------------------")
print("Input Dimension: {}".format(tree_labels.shape[1]))
print("Output Dimension: {}".format(tree_features.shape[1]))
print("Number of samples: {}".format(tree_features.shape[0]))
```

```python
def exercise_1(η=100.0, epochs=250 , γ=1e-12):
  """Training a LNN

  Args:
  η (float): learning rate (default 100.0)
  epochs (int): number of epochs (default 250)
  γ (float): initialization scale (default 1e-12)

  """
  n_hidden = [30]

  dim_input = tree_labels.shape[1]
  dim_output = tree_features.shape[1]

  deep_model = VariableDepthWidth(in_dim=dim_input,
                                     out_dim=dim_output,
                                     hid_dims=n_hidden,
                                     gamma=γ)

  # convert (cast) data from np.ndarray to torch.Tensor
  input_tensor = torch.tensor(tree_labels).float()
  #################################################
  ## convert output_data from np.ndarray to torch.Tensor
  # Complete the function and remove or comment the line below
  raise NotImplementedError("Cast output_data as torch.Tensor")
  #################################################
  output_tensor = ...

  training_losses = train(deep_model,
                          input_tensor,
                          output_tensor,
                          n_epochs=epochs,
                          lr=η)

  plot_loss(training_losses)

# # Uncomment and run
# exercise_1()
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content-dl/tree/main//tutorials/W1D2_LinearNN/solutions/W1D2_Tutorial3_Solution_2c03a4ed.py)

*Example output:*

**Question**: Why haven't we seen these "bumps" in training before? And should we look for them in the future? You can slide the widgets below and find your answer. Here, $\gamma$ is the initialization scale.

```python
#@markdown Make sure you execute this cell to enable the widget!
_ = interact(exercise_1,
             η = FloatSlider(min=1.0, max=200.0, step=2.0, value=100.0,
                             continuous_update=False, readout_format='.1f', description='η'),
             epochs = fixed(250),
             γ = FloatLogSlider(min=-15, max=1, step=1, value=1e-12, base=10,
                             continuous_update=False, description='γ'),
             )
```

---
# Section 2: Singular Value Decomposition (SVD)

```python
#@title Video 2: Singular Value Decomposition (SVD)
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="eTXNKMleEj8", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

In this section, we would go deeper in understanding of the learning dynamics we just saw. First, we should know that a linear neural network is perfoming sequential matrix maltiplications, which can simplified to:

\begin{align}
\mathbf{y} &= \mathbf{W}_{L}~\mathbf{W}_{L-1}~\dots~\mathbf{W}_{1} ~ \mathbf{x} \\
 &= (\prod_{i=1}^{L}{\mathbf{W}_{i}}) ~ \mathbf{x} \\
 &= \mathbf{W}_{tot} ~ \mathbf{x}
\end{align}

where $L$ denotes the number of layers in our network.

Why did we just called the learning progress, "learning dynamics"? Learning through gradient descent seems very aike to evolution of a dynamic system. They both are described by set of differntial equations (gradients). Dynamical systems often have a "time-constant" which describes the rate of change, similar to the learning rate, only instead of time, gradient descent evolves through epochs.

[Saxe et al. (2013)](https://arxiv.org/abs/1312.6120) showed that to analyse and to understanding the nonlinear learning dynamics of a deep LNN, we can use [Singular Value Decomposition (SVD)](https://en.wikipedia.org/wiki/Singular_value_decomposition) to decompose the $\mathbf{W}_{tot}$ into orthogonal vectors, where orthogonality of the vecors would ensure their "individuality". This means we can break the a deep wide LNN to multiple deep narrow LNN, so their activity is untangled from each other.

__A Quick intro to SVD__

Any real-valued matix $A$ (yes, ANY) can be decomposed (factorized) to 3 matrices:

\begin{equation}
\mathbf{A} = \mathbf{U} \mathbf{Σ} \mathbf{V}^{\top}
\end{equation}

where $U$ is an orthogonal matrix, $\Sigma$ is a diagonal matrix, and $V$ is again an orthogonal matrix. The diagonal elements of $\Sigma$ are called **singular values**.

The main difference between SVD and Eigen Value Decomposition (EVD), is that EVD requires $A$ to be squred and does not guarantee the eigenvectors to be orthogonal. For the complex-valued matrix $A$, the factorization changes to $A = UΣV^*$ and $U$ and $V$ are unitary matrices.

We strongly recommend the [Singular Value Decomposition (the SVD)](https://www.youtube.com/watch?v=mBcLRGuAFUk) by the amazing [Gilbert Strang](http://www-math.mit.edu/~gs/) if you would like to learn more.

## Exercise 2. SVD

Let's put what we learned in practice. Here, we want to modify our training loop to perform the SVD on $\mathbf{W}_{tot}$ in every epoch, and record the singular values (diagonal values of $\Sigma$). SVD is implemented both in PyTorch [`torch.svd`](https://pytorch.org/docs/stable/generated/torch.svd.html) and in NumPy [`np.linalg.svd`](https://numpy.org/doc/stable/reference/generated/numpy.linalg.svd.html), but we recommend the PyTorch method to avoid the conversion cost. Since $\Sigma$ is a diagonal matrix, often (e.g. PyTorch and NumPy) just the diagonal elements are returned as a vector, not the whole matrix.

We have removed the progress bar and the optional loss and optimizer arguments to make the exercise "cleaner".

```python
def train_svd_exercise(model, in_features, out_features, n_epochs, lr):
  """Training function

  Args:
    model (torch nn.Module): the neural network
    in_features (torch.Tensor): features (input) with shape `torch.Size([batch_size, input_dim])`
    out_features (torch.Tensor): targets (labels) with shape `torch.Size([batch_size, output_dim])`
    n_epochs (int): number of training epochs
    lr(float): learning rate

  Returns:
    np.ndarray: record (evolution) of losses
    np.ndarray: record (evolution) of singular values
  """

  assert in_features.shape[0] == out_features.shape[0]
  optimizer = optim.SGD(model.parameters(), lr=lr)
  criterion = nn.MSELoss()
  loss_record = []  # for recoding losses
  sv_record = []  # for recoding singular values

  for i in range(n_epochs):
    y_pred = model(in_features)  # forward pass
    loss = criterion(y_pred, out_features)  # calculating the loss
    optimizer.zero_grad()  # reset all the graph gradients to zero
    loss.backward()  # back propagation of the error
    optimizer.step()  # gradient step

    # calculating the W_tot by multiplying all layers' weights
    W_tot = model.layers[-1].weight.detach()  # starting from the last layer
    for i in range(2, len(model.layers)+1):
      #################################################
      ## Complete the loop for calculating the W_tot
      # Complete the function and remove or comment the line below
      raise NotImplementedError("Calculate the W_tot")
      #################################################
      W_tot = ...

    # performing the SVD!
    #################################################
    ## calculate singular value decomposition of W_tot
    # Complete the function and remove or comment the line below
    raise NotImplementedError("Calculate the SVD for W_tot")
    #################################################
    U, Σ, V = ...

    loss_record.append(loss.item())
    sv_record.append(Σ.numpy())

  return np.array(loss_record), np.array(sv_record)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content-dl/tree/main//tutorials/W1D2_LinearNN/solutions/W1D2_Tutorial3_Solution_047b33fe.py)

```python
#@markdown Make sure you execute this cell to train the network and plot

dim_input = tree_labels.shape[1]
dim_output = tree_features.shape[1]
input_tensor = torch.tensor(tree_labels).float()
output_tensor = torch.tensor(tree_features).float()

deep_model = VariableDepthWidth(in_dim=dim_input,
                                    out_dim=dim_output,
                                    hid_dims=[30])

training_losses, singular_values, _, _ = train_svd_rsa_track(deep_model,
                        input_tensor,
                        output_tensor,
                        n_epochs=250,
                        lr=100.0)

plot_loss_sv_twin(training_losses, singular_values)
```

**Question**: Isn't this beautiful? For EigenValue decomposition, the anount of variance explained by eigenvectors are proportional to the corresponding eigenvalues. What about the SVD? We definitely see that the gradient descent guides the network to first learn the features that carry more information (have higher singular value)!

---
# Section 3: Representational Similarity Analysis (RSA)

```python
#@title Video 3.1: Representational Similarity Analysis (RSA)
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="19seHV97WkI", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

The previous section ended with an interesting remark! The network (through gradient descent), seems to prioritize learning features that explain most of the data, and gradually learn all the hidden representations. Given that we are training a hierarchically structured data, we may be able to see that progress as well.

To do so, we get help from Representational Similarity Analysis (RSA) approach to understand the internal representation of our network. The main idea is that the activity of hidden units (neurons) in the network must be similar when the network is presented with a similar input. The exercise will help more to get an intuition of this approach and the dynamics of representational learning.

## Exercise 3: RSA

We need to modify our training function once more. The task is to calculate similarity between the hidden layer activities (i.e. $~\mathbf{h_1} = \mathbf{W_1} \mathbf{x}~$) for all the inputs at every epoch. For similarity measure, we can use the good old dot (scalar) product, which is also called cosine similarity. For calculating the dot product between multiple vectors (which would be our case), you can simply use matrix multiplication. Therefore the Representational Similarity Matrix (SM) for multiple input activity could be calculated as follow:

$$ RSM = \mathbf{H_1}^{\top} \mathbf{H_1} $$

where $\mathbf{H_1} = \mathbf{W_1} \mathbf{X}$.

```python
def train_svd_rsa_exercise(model, in_features, out_features, n_epochs, lr):
  """Training function

  Args:
    model (torch nn.Module): the neural network
    in_features (torch.Tensor): features (input) with shape `torch.Size([batch_size, input_dim])`
    out_features (torch.Tensor): targets (labels) with shape `torch.Size([batch_size, output_dim])`
    n_epochs (int): number of training epochs
    lr(float): learning rate

  Returns:
    np.ndarray: record (evolution) of losses
    np.ndarray: record (evolution) of singular values
    np.ndarray: record (evolution) of representational similarity matrices
  """

  assert in_features.shape[0] == out_features.shape[0]
  optimizer = optim.SGD(model.parameters(), lr=lr)
  criterion = nn.MSELoss()
  loss_record = []  # for recoding losses
  sv_record = []  # for recoding singular values
  rsm_record = []  # for recording representational similarity matrices

  for i in range(n_epochs):
    y_pred = model(in_features)  # forward pass
    loss = criterion(y_pred, out_features)  # calculating the loss
    optimizer.zero_grad()  # reset all the graph gradients to zero
    loss.backward()  # back propagation of the error
    optimizer.step()  # gradient step

    # calculating the W_tot by multiplying all layers' weights
    W_tot = model.layers[-1].weight.detach()  # starting from the last layer
    for i in range(2, len(model.layers)+1):
      W_tot = W_tot @ model.layers[-i].weight.detach()
    U, Σ, V = torch.svd(W_tot)  # performing the SVD!

    # calculating representational similarity matrix
    H1 = model.layers[0].weight.detach() @ in_features
    #################################################
    ## Use H1 to calculate the representational similarity matrix
    # Complete the function and remove or comment the line below
    raise NotImplementedError("Calculate the RSM")
    #################################################
    RSM = ...

    loss_record.append(loss.item())
    sv_record.append(Σ.numpy())
    rsm_record.append(RSM.numpy())

  return np.array(loss_record), np.array(sv_record), np.array(rsm_record)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content-dl/tree/main//tutorials/W1D2_LinearNN/solutions/W1D2_Tutorial3_Solution_78a1d0b5.py)

```python
#@markdown Make sure you execute this cell to train the network

deep_model = VariableDepthWidth(in_dim=dim_input,
                                    out_dim=dim_output,
                                    hid_dims=[30])

training_losses, singular_values, rep_sim_mats, _ = train_svd_rsa_track(deep_model,
                        input_tensor,
                        output_tensor,
                        n_epochs=250,
                        lr=100.0)
```

Using the widget below, you can look at the representational similartiy matrix at any point of training.

```python
#@markdown Make sure you execute this cell to enable the widget!

i_ep_slider = IntSlider(min=5, max=245, step=1, value=50,
                        continuous_update=False, description='Epoch',
                        layout=Layout(width='680px'))

widgets_ui = HBox([i_ep_slider])

widgets_out = interactive_output(plot_loss_sv_rsm,
                                         {'loss_array': fixed(training_losses),
                                          'sv_array': fixed(singular_values),
                                          'rsm_array': fixed(rep_sim_mats),
                                          'i_ep': i_ep_slider})

display(widgets_ui, widgets_out)
```

```python
#@title Video 3.2: Linear Regression
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="etsXyJJSru4", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

## Demonstration: Linear Regression vs. DLNN

A linear neural network with NO hidden layer, is very similar to linear regression in its core. We also know that no matter how many hidden layers a linear network has, it can be compressed to a linear regression (no hidden layers).

In this demonstration, we use the hierarchically structured data to:

* analytically find the mapping between features and labels
* train a zero-depth LNN to find the mapping
* compare them to the $W_{tot}$ from the already trained deep LNN

```python
# calculating the W_tot for deep network (already trained model)
deep_weight_tot = deep_model.layers[-1].weight.detach().numpy()
for i in range(2, len(deep_model.layers)+1):
  deep_weight_tot = deep_weight_tot @ deep_model.layers[-i].weight.detach().numpy()
```

```python
# analytically estimation of weights (map)
# our data is batch first dimension, so we need to transpose our data
analytical_weights = linear_regression(tree_labels.T, tree_features.T)
```

```python
# create a model instance of VariableDepthWidth
zero_depth_model = VariableDepthWidth(in_dim=dim_input,
                                    out_dim=dim_output,
                                    hid_dims=[])
# train the zero_depth_model
training_losses = train(zero_depth_model,
                        input_tensor,
                        output_tensor,
                        n_epochs=250,
                        lr=1000.0)
# trained weights from zero_depth_model
zero_depth_model_weights = zero_depth_model.layers[0].weight.detach().numpy()

plot_loss(training_losses, "Training loss for zero depth LNN", c="r")
```

```python
print("The final weights from all methods are approximately equal?! {}!\n".format(
  (np.allclose(analytical_weights, zero_depth_model_weights, atol=1e-02) and \
   np.allclose(analytical_weights, deep_weight_tot, atol=1e-02))))
```

As you may have guessed, they all arrive at the same results but through very different paths.

---
# Section 4: Illusory Correlations

```python
#@title Video 4.1: Illusory Correlations
# Insert the ID of the corresponding youtube video
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="t_-wmMjl9kk", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

So far, everything looks great, all our trainings are successful (training loss converging to zero), and very fast. We even could interpret the dynamics of our deep linear networks and relate them to the data. Unfortunately, this rarely happens in practice. Real world problems often require very deep and nonlinear networks with many hyper-parameters. And oridinarily, these complex networks take hours, if not days, to train.

Let's recall the training loss curves. There was often a long plateau (where the weights are stuck at a saddle point), followed by a sudden drop. For very deep complex neural network, such plateaus can last for hours of training, and we often decide to stop the training becasue we believe it "as good as it gets"! This riases a challenge of whether the network has learned all the "intended" hidden representations. But more importantly, the network might find an illusionary correlation between features that has never seen.

To better understand this, let's do the next demonstration and exercise.

## Demonstration: Illusory Correlations

So far we worked with a dataset that has 4 animals: Canary, Robin, Goldfish, and Tuna. These animals all have bones. Therefore if we include the "has bone" feature, the network would learn it at the second level (i.e. second bump, second singular value convergence), which is OK.

What if the dataset has Shark instead of Goldfish. Sharks don't have bones (their skeletons is made of cartilaginous, which is much lighter than true bone and more flexible). Then we will have feature which is *True* (i.e. +1) for Tuna, Robin, and Canary, but *False* (i.e. 0) for all the plants and the shark! Let's see what the network does.

First, we add the new feature to the targets. We then start training our LNN and in every epoch, record the network prediction for "sharks having bones".

```python
item_names = ['Shark', 'Tuna', 'Robin', 'Canary', 'Rose', 'Daisy', 'Pine', 'Oak']
has_bones = [0, 1, 1, 1, 0, 0, 0, 0]

tree_features = add_feature(tree_features, has_bones)
plot_tree_data(tree_labels, tree_features, item_names)
```

You can see the new feature shown in the last column of the plot above.

```python
#@markdown Make sure you execute this cell to train the network
dim_input = tree_labels.shape[1]
dim_output = tree_features.shape[1]
input_tensor = torch.tensor(tree_labels).float()
output_tensor = torch.tensor(tree_features).float()

deep_model = VariableDepthWidth(in_dim=dim_input,
                                out_dim=dim_output,
                                hid_dims=[30])

_, singular_values, _, ill_predictions = train_svd_rsa_track(deep_model,
                                              input_tensor,
                                              output_tensor,
                                              n_epochs=250,
                                              lr=100.0,
                                              ill_i=0)

plot_ills_sv_twin(ill_predictions, singular_values)
```

It seems that the network starts by learning the "[alternative fact](https://en.wikipedia.org/wiki/Alternative_facts)" that sharks have bones, and in later epochs, as it learns deeper representations, it can see (learn) beyond the illusory correlation. This is important to remember that we never presented the network with any data saying that sharks have bones.

```python
#@title Video 4.2: Illusory Correlations Explained
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="QMuTlq-atlc", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

## Exercise 4: Illusory Correlations

This exercise is just for you to explore the idea of illusory correlations. Think of medical, natural or possibly social illusory correlations which can test the learning power of deep linear neural nets.

**Notes**: Before you start, there are few things to important to know:

* the generated data is independent of tree labels, therefore the names are just for convenience.

* you can rename any node in the tree object(from the `SimpleTree` class) to help you keep track of your tree. This tree, also is not the one generating the samples and is purely for convenience.

Here is our example for **Non-human Living things don't speak**:

```python
# this is just for plotting a tree and has no connection to data!
tree = SimpleTree()  # creates a tree
tree.rename("Canary", "Parrot")  # renames the Canary node to Parrot
tree.plot()  # plots the tree
```

```python
item_names = ['Goldfish', 'Tuna', 'Robin', 'Parrot', 'Rose', 'Daisy', 'Pine', 'Oak']
can_NOT_speak = [1, 1, 1, 0, 1, 1, 1, 1]  # creating the new feature
ill_id = 3  # the index of your feature

tree_labels, tree_features = generate_hsd()  # sampling new data from the tree
tree_features = add_feature(tree_features, can_NOT_speak) # adding the feature
plot_tree_data(tree_labels, tree_features, item_names)  # plot
```

```python
#@markdown Make sure you execute this cell to train the network and plot the output
dim_input = tree_labels.shape[1]
dim_output = tree_features.shape[1]
input_tensor = torch.tensor(tree_labels).float()
output_tensor = torch.tensor(tree_features).float()

deep_model = VariableDepthWidth(in_dim=dim_input,
                                out_dim=dim_output,
                                hid_dims=[30])

_, singular_values, _, ill_predictions = train_svd_rsa_track(deep_model,
                                              input_tensor,
                                              output_tensor,
                                              n_epochs=250,
                                              lr=100.0,
                                              ill_i=ill_id)

plot_ills_sv_twin(ill_predictions, singular_values)
```

---
# Wrap up

```python
#@title Video 4.3: Outro
# Insert the ID of the corresponding youtube video
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="Y0JfyCtikhc", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

---
# Appendix

Generally, *regression* refers to a set of methods for modeling the mapping (relationship) between one (or more) independent variable(s) (i.e. features) and one (or more) dependent variable(s) (i.e. labels). For example, if we want to examine the relative impacts of calendar date, GPS coordinates, and time of the say (the independent variables) on air temperature (the dependent variable). On the other hand, regression can be used for predictive analysis. Thus the independent variables are also called predictors. When the model contains more than one predictor, then the method is called *multiple regression*, and if it contains more than one dependent variable called *multivariate regression*. Regression problems pop up whenever we want to predict a numerical (usually continuous) value.

The independent variables are collected in vector $\mathbf{x} \in \mathbb{R}^M$, where $M$ denotes the number of independent variables, while the dependent variables are collected in vector $\mathbf{y} \in \mathbb{R}^N$, where $N$ denotes the number of independent variables. And the mapping between them is represented by the weight matrix $\mathbf{W} \in \mathbb{R}^{N \times M}$ and a bias vector $\mathbf{b} \in \mathbb{R}^{N}$ (generalizing to affine mappings).

The multivariate regression model can be written as:

\begin{equation}
\mathbf{y} = \mathbf{W} ~ \mathbf{x} + \mathbf{b}
\end{equation}

or it can be written in matrix format as:

\begin{equation}
\begin{bmatrix} y_{1} \\ y_{2} \\ \vdots \\ y_{N} \\ \end{bmatrix} = \begin{bmatrix} w_{1,1} & w_{1,2} & \dots & w_{1,M} \\ w_{2,1} & w_{2,2} & \dots & w_{2,M} \\ \vdots & \ddots & \ddots & \vdots \\ w_{N,1} & w_{N,2} & \dots & w_{N,M} \end{bmatrix} \begin{bmatrix} x_{1} \\ x_{2} \\ \vdots \\ x_{M} \\ \end{bmatrix} + \begin{bmatrix} b_{1} \\ b_{2} \\ \vdots \\b_{N} \\ \end{bmatrix}
\end{equation}

__Vectorized regression__

Linear regression can be simply extended to multi-samples ($D$) input-output mapping, which we can collect in a matrix $\mathbf{X} \in \mathbb{R}^{M \times D}$, sometimes called the design matrix. The sample dimension also shows up in the output matrix $\mathbf{Y} \in \mathbb{R}^{N \times D}$. Thus, linear regression takes the following form:

\begin{equation}
\mathbf{Y} = \mathbf{W} ~ \mathbf{X} + \mathbf{b}
\end{equation}

where matrix $\mathbf{W} \in \mathbb{R}^{N \times M}$ and the vector $\mathbf{b} \in \mathbb{R}^{N}$ (broudcasted over sample dimension) are the desired parameters to find.