<jupyter_start><jupyter_text>Simple shape recognition<jupyter_code>import lettertask as lt
import patches
import torch
import numpy as np
import torch.optim as optim
import torch.nn as nn
import pandas as pd
import plotnine as gg
from tqdm import tqdm
import lazytools_sflippl as lazytools
agbm = lt.data.AtomicGeometricBinaryModel(width=2, samples=100000)
agbm_flat = agbm.to_array().reshape(-1, 4)
agbm_flat<jupyter_output><empty_output><jupyter_text>Supervised model<jupyter_code>cts = patches.data.HiddenMarkovModel(agbm_flat, agbm.values)
class SupervisedCNN(nn.Module):
    def __init__(self, latent_features, input_features=None,
                 data=None, bias=True):
        super().__init__()
        if data:
            input_features = input_features or data.n_vars
        elif input_features is None or timesteps is None:
            raise ValueError('You must either provide data or both input '
                             'features and timesteps.')
        self.latent_features = latent_features
        self.recognition = nn.Linear(input_features, 2*latent_features)
        self.nonlinearity = nn.ReLU()
        self.shifter = nn.Linear(2*latent_features, latent_features)

    def forward(self, x):
        latent = self.recognition(x['input'])
        latent = self.nonlinearity(latent)
        latent = self.shifter(latent)
        return latent
scnn = SupervisedCNN(1, data=cts)
cdl = torch.utils.data.DataLoader(cts, batch_size=8, drop_last=True)
optimizer = optim.Adam(scnn.parameters())
criterion = nn.MSELoss()
loss_traj = []
running_loss = 0
for epoch in tqdm(range(10)):
    running_loss = 0
    for i, data in enumerate(cdl):
        if i<len(cdl):
            optimizer.zero_grad()
            prediction = scnn(data)
            loss = criterion(prediction, data['latent_values'])
            loss.backward()
            optimizer.step()
            running_loss += loss
            if i % 50 == 49:
                loss_traj.append(running_loss.detach().numpy()/50)
                running_loss = 0
(gg.ggplot(lazytools.array_to_dataframe(np.array(loss_traj)),
           gg.aes(x='dim0', y='array')) +
    gg.geom_line(alpha=0.2) +
    gg.geom_smooth())
list(scnn.parameters())
def reaction(module):
    shapes = [
        torch.tensor([[1., 1], [0, 0]]),
        torch.tensor([[0., 0], [1, 1]]),
        torch.tensor([[1., 0], [1, 0]]),
        torch.tensor([[0., 1], [0, 1]])
    ]
    reaction = []
    for arr in shapes:
        reaction.append(module.nonlinearity((module.recognition(arr.flatten()))).detach().numpy())
    return np.array(reaction)
reaction(scnn)
list(scnn.parameters())<jupyter_output><empty_output><jupyter_text>Supervised prediction<jupyter_code>class SupervisedPredictiveCNN(nn.Module):
    def __init__(self, latent_features, input_features=None,
                 data=None, bias=True):
        super().__init__()
        if data:
            input_features = input_features or data.n_vars
        elif input_features is None or timesteps is None:
            raise ValueError('You must either provide data or both input '
                             'features and timesteps.')
        self.latent_features = latent_features
        self.recognition = nn.Linear(input_features, 2*latent_features)
        self.nonlinearity = nn.ReLU()
        self.shifter = nn.Linear(2*latent_features, latent_features)
        self.predictor = nn.Linear(latent_features, 5*latent_features)

    def forward(self, x):
        latent = self.recognition(x['input'])
        latent = self.nonlinearity(latent)
        latent = self.shifter(latent)
        latent = self.predictor(latent).reshape(-1, 5, self.latent_features)
        return latent
spcnn = SupervisedPredictiveCNN(1, data=cts)
optimizer = optim.Adam(spcnn.parameters())
criterion = nn.MSELoss()
loss_traj = []
running_loss = 0
for epoch in tqdm(range(10)):
    running_loss = 0
    for i, data in enumerate(cdl):
        if i<len(cdl):
            optimizer.zero_grad()
            prediction = spcnn(data)
            loss = criterion(prediction, data['future_latent_values'])
            loss.backward()
            optimizer.step()
            running_loss += loss
            if i % 50 == 49:
                loss_traj.append(running_loss.detach().numpy()/50)
                running_loss = 0
(gg.ggplot(lazytools.array_to_dataframe(np.array(loss_traj)),
           gg.aes(x='dim0', y='array')) +
    gg.geom_line(alpha=0.2) +
    gg.geom_smooth())
reaction(spcnn)<jupyter_output><empty_output><jupyter_text>Contrastive coding<jupyter_code>ccts = patches.data.Contrastive1DTimeSeries(agbm_flat, n_contrasts=10)
class SupervisedCNN(nn.Module):
    def __init__(self, latent_features, input_features=None,
                 data=None, bias=True):
        super().__init__()
        if data:
            input_features = input_features or data.n_vars
        elif input_features is None or timesteps is None:
            raise ValueError('You must either provide data or both input '
                             'features and timesteps.')
        self.latent_features = latent_features
        self.recognition = nn.Linear(input_features, 10*latent_features)
        self.nonlinearity = nn.ReLU()
        self.layer_2 = nn.Linear(10*latent_features, 2*latent_features)
        self.shifter = nn.Linear(2*latent_features, latent_features)

    def forward(self, x):
        latent = self.recognition(x)
        latent = self.nonlinearity(latent)
        latent = self.layer_2(latent)
        latent = self.shifter(latent)
        return latent
class TwoLayeredScaffold(patches.networks.SimpleScaffold):
    def __init__(self):
        super().__init__()
        self.encoder = SupervisedCNN(1, data=ccts)
        self.predictor = nn.Linear(1, 5)

    def encode(self, x):
        return self.encoder(x)

    def predict(self, x):
        return self.predictor(x)
tls = TwoLayeredScaffold()
criterion = patches.losses.ContrastiveLoss(loss=patches.losses.BilinearLoss())
cdl = torch.utils.data.DataLoader(ccts, drop_last=True, batch_size=8)
optimizer = optim.Adam(tls.parameters(), lr=2e-4)
loss_traj = []
running_loss = 0
n_epochs = 2
with tqdm(total=n_epochs*len(cdl)) as pbar:
    for epoch in range(n_epochs):
        running_loss = 0
        for i, data in enumerate(cdl):
            if i<len(cdl):
                optimizer.zero_grad()
                code = tls(data)
                loss = criterion(code)
                loss.backward()
                optimizer.step()
                running_loss += loss
                if i % 50 == 49:
                    loss_traj.append(running_loss.detach().numpy()/50)
                    running_loss = 0
                pbar.update(1)
(gg.ggplot(lazytools.array_to_dataframe(np.array(loss_traj)),
           gg.aes(x='dim0', y='array')) +
    gg.geom_line(alpha=0.2) +
    gg.geom_smooth())
reaction(tls.encoder)
tls.encode(torch.tensor([[1., 1, 0, 0],
                 [0, 0, 1, 1],
                 [1, 0, 1, 0],
                 [0, 1, 0, 1]]))<jupyter_output><empty_output><jupyter_text>For such a complex model, there is a sudden shift after four iterations -- and we have found the right dimension. This suggests a highly nonlinear loss landscape and both the linear and the nonlinear case might be suitable for a more principled mathematical investigation of this landscape as it is an important determinant for the algorithm's performance at more complex tasks. For now, we attempt to distill the model.For the single nonlinearity, many epochs seemed to not change anything. In light of this new evidence, we will now revisit this case.<jupyter_code>class SupervisedCNN(nn.Module):
    def __init__(self, latent_features, input_features=None,
                 data=None, bias=True):
        super().__init__()
        if data:
            input_features = input_features or data.n_vars
        elif input_features is None or timesteps is None:
            raise ValueError('You must either provide data or both input '
                             'features and timesteps.')
        self.latent_features = latent_features
        self.recognition = nn.Linear(input_features, 2*latent_features)
        self.nonlinearity = nn.ReLU()
        self.shifter = nn.Linear(2*latent_features, latent_features)

    def forward(self, x):
        latent = self.recognition(x)
        latent = self.nonlinearity(latent)
        latent = self.shifter(latent)
        return latent
tls = TwoLayeredScaffold()
optimizer = optim.Adam(tls.parameters(), lr=2e-3)
loss_traj = []
running_loss = 0
n_epochs = 60
with tqdm(total=n_epochs*len(cdl)) as pbar:
    for epoch in range(n_epochs):
        running_loss = 0
        for i, data in enumerate(cdl):
            if i<len(cdl):
                optimizer.zero_grad()
                code = tls(data)
                loss = criterion(code)
                loss.backward()
                optimizer.step()
                running_loss += loss
                if i % 50 == 49:
                    loss_traj.append(running_loss.detach().numpy()/50)
                    running_loss = 0
                pbar.update(1)
(gg.ggplot(lazytools.array_to_dataframe(np.array(loss_traj)),
           gg.aes(x='dim0', y='array')) +
    gg.geom_line(alpha=0.2))
-np.log(1/11)
np.array(loss_traj).mean()
tls.encode(torch.tensor([[1., 1, 0, 0],
                 [0, 0, 1, 1],
                 [1, 0, 1, 0],
                 [0, 1, 0, 1]]))<jupyter_output><empty_output><jupyter_text>We have therefore failed to find one dimension. What if we add one more feature dimension?<jupyter_code>class SupervisedCNN(nn.Module):
    def __init__(self, latent_features, input_features=None,
                 data=None, bias=True):
        super().__init__()
        if data:
            input_features = input_features or data.n_vars
        elif input_features is None or timesteps is None:
            raise ValueError('You must either provide data or both input '
                             'features and timesteps.')
        self.latent_features = latent_features
        self.recognition = nn.Linear(input_features, 3*latent_features)
        self.nonlinearity = nn.ReLU()
        self.shifter = nn.Linear(3*latent_features, latent_features)

    def forward(self, x):
        latent = self.recognition(x)
        latent = self.nonlinearity(latent)
        latent = self.shifter(latent)
        return latent
tls = TwoLayeredScaffold()
optimizer = optim.Adam(tls.parameters(), lr=2e-3)
loss_traj = []
running_loss = 0
n_epochs = 60
with tqdm(total=n_epochs*len(cdl)) as pbar:
    for epoch in range(n_epochs):
        running_loss = 0
        for i, data in enumerate(cdl):
            if i<len(cdl):
                optimizer.zero_grad()
                code = tls(data)
                loss = criterion(code)
                loss.backward()
                optimizer.step()
                running_loss += loss
                if i % 50 == 49:
                    loss_traj.append(running_loss.detach().numpy()/50)
                    running_loss = 0
                pbar.update(1)
(gg.ggplot(lazytools.array_to_dataframe(np.array(loss_traj)),
           gg.aes(x='dim0', y='array')) +
    gg.geom_line(alpha=0.2))
tls.encode(torch.tensor([[1., 1, 0, 0],
                 [0, 0, 1, 1],
                 [1, 0, 1, 0],
                 [0, 1, 0, 1]]))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11098398169336385, 'fraction_numerical': 0.016106319309980636, 'mean_word_length': 2.4952322362350046, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '330705', 'license': 'MIT', 'n_tokens_mistral': 3813, 'n_tokens_neox': 3578, 'n_words': 794, 'path': 'notebooks/15-gbm-simplified.ipynb', 'repo_name': 'sflippl/patches', 'size': '11362'}