In [ ]:
import lsrl
from tests.test_ua_vec2vec import TestUAVec2Vec
from tests.test_ua_tok2tok import TestUATok2TokSimplification, TestUATok2TokBuldingSteps
import copy
import numpy as np
from collections import defaultdict
import tqdm.notebook as tqdm
from matplotlib import pyplot as plt
import pandas
import seaborn as sns
from matplotlib.colors import LogNorm
from joblib import Parallel, delayed
from itertools import product
Test the effect of parameter noise on the Vec2Vec model¶
In [ ]:
overall_results = defaultdict(list)
n_targets = 25
def process_one_target(random_seed: int, noise_level: float, conditional_mode: str):
# set the random seed
np.random.seed(random_seed)
# creates a target function to estimate
target = lsrl.utils.create_test_target_function(
in_dim=2,
out_dim=2,
n_control_points=5,
smoothness=1,
)
discretization_levels = 8
# creates the prompt describing the target function
prompts = lsrl.utils.construct_prompt(
target,
in_dim=2,
out_dim=2,
N=discretization_levels,
)
# random samples from the hypercube where we evaluate the quality of the approximation
test_x_vals = np.random.random((10, 2))
test_y_vals = np.array([target(x) for x in test_x_vals])
inner_results = dict()
# PREPARE THE ORIGINAL (UNNOISED MODEL)
tmp_test_constructor = TestUAVec2Vec()
tmp_test_constructor.conditional_mode = conditional_mode
tmp_test_constructor.mode = lsrl.MatrixMode.NUMERIC
tmp_test_constructor.make_plots = False
tmp_test_constructor.setUp()
model_original = copy.deepcopy(tmp_test_constructor.simplified)
# this hook records the states during the unnoised pass
# we want to input the unnoised states at each step rather than the noisy outputs
# that's because the errors will otherwise blow up just because of the larger number of steps
# and this prevents us to conclude whether that's due to the f_ifelse implementation
# this way these results are independent of the prompt length
states = defaultdict(list)
def record_state_hook(s, model, self, *args, **kwargs):
if isinstance(self, lsrl.LinState):
node_idx = model.topological_sort().index(self)
s[node_idx].append(copy.deepcopy(self._state))
# add the hook that records outputs to all nodes
for node in model_original.graph().nodes():
node.hooks["before_forward"].append(lambda self, *args, **kwargs: record_state_hook(states, model_original, self, *args, **kwargs))
# PREPARE THE NOISY MODEL
model_noisy = copy.deepcopy(tmp_test_constructor.simplified)
# go through the nodes and each node's Matrix type attribute and add a little gaussian noise to it
# remember that the Matrices are scipy sparse matrices
if noise_level != 0:
for node in model_noisy.graph().nodes():
for attr in node.__dict__:
if isinstance(getattr(node, attr), lsrl.Matrix) and attr in ["A", "B", "b", "bias"]:
new_data = getattr(node, attr).matrix.data + np.random.normal(0, noise_level, getattr(node, attr).matrix.data.shape)
new_matrix = copy.deepcopy(getattr(node, attr))
new_matrix.matrix.data = new_data
setattr(node, attr, new_matrix)
# add a hook to the LinStates setting the state to the unnoised one.
# That's because we want to evaluate the effect of noising without complicating
# the analysis with error accumulation through the states. If didn't do that,
# the error rate would be monotonically increasing in the prompt length
def fix_state_hook(model, self, *args, **kwargs):
if isinstance(self, lsrl.LinState):
node_idx = model.topological_sort().index(self)
original_state = states[node_idx].pop(0)
# print(f"replacing state {self._state.numpy().flatten()} with {original_state.numpy().flatten()}")
self._state = original_state
for node in model_noisy.graph().nodes():
if isinstance(node, lsrl.LinState):
node.hooks["before_forward"].append(lambda self, *args, **kwargs: fix_state_hook(model_noisy, self, *args, **kwargs))
# EVALUATE AT ALL INPUTS
results_original = []
results_noisy = []
for x_val in test_x_vals:
first_input = np.concatenate((x_val, np.zeros(tmp_test_constructor.target_fun_output_dim + 1)))
full_inputs = np.vstack((first_input, prompts))
states.clear()
results_original.append(model_original(full_inputs.T).numpy()[:, -1].flatten())
results_noisy.append(model_noisy(full_inputs.T).numpy()[:, -1].flatten())
error_noisy = np.linalg.norm(np.array(results_noisy)-test_y_vals, axis=1)
inner_results[(conditional_mode, noise_level)] = np.mean(error_noisy)
return inner_results
# run n_targets instances in parallel and combine their results
modes = ["multiplicative", "direct", "optimized", "step_optimized", "multiplicative_optimized"]
results = Parallel(n_jobs=16)(delayed(process_one_target)(i, noise, mode) for i, noise, mode in tqdm.tqdm(list(product(range(n_targets), [0, 1e-3, 1e-5, 1e-7, 1e-9], modes))))
for res in results:
for k, v in res.items():
overall_results[k].append(v)
df = pandas.DataFrame({"conditional_mode": m, "Noise level": n, "Average difference": e} for (m, n), l in overall_results.items() for e in l)
df.to_csv("noise_robustness_vec2vec.csv")
In [ ]:
df = pandas.read_csv("noise_robustness_vec2vec.csv")
sns.stripplot(
hue="conditional_mode",
y="Average difference",
x="Noise level",
data=df,
log_scale=True,
dodge=True,
hue_order=["direct", "optimized", "step_optimized", "multiplicative", "multiplicative_optimized"],
alpha=0.5,
)
plt.grid(axis="y", which="both")
plt.savefig("noise_effect_vec2vec.pdf", bbox_inches="tight")
Test the effect of parameter noise on the Tok2Tok model¶
In [ ]:
overall_results = defaultdict(list)
n_targets = 25
def process_one_target(random_seed: int, noise_level: float, conditional_mode: str):
# set the random seed
np.random.seed(random_seed)
# create a random map for us to model
tmp = TestUATok2TokBuldingSteps()
tmp.mode = lsrl.MatrixMode.NUMERIC
tmp.conditional_mode = "direct"
tmp.seq_len = 10
tmp.n_samples = 20
tmp.setUp()
expected = []
full_inputs = []
inner_results = dict()
# generate test samples
while len(expected) < 5:
x, y, full_input = tmp.prep_random_input()
# prevent repetitions
if any([(x == fi[: len(x)]).all() for fi in full_inputs]):
continue
full_inputs.append(full_input)
expected.append(y)
# PREPARE THE ORIGINAL (UNNOISED MODEL)
tmp_test_constructor = TestUATok2TokSimplification()
tmp_test_constructor.conditional_mode = conditional_mode
tmp_test_constructor.mode = lsrl.MatrixMode.NUMERIC
tmp_test_constructor.make_plots = False
tmp_test_constructor.seq_len = tmp.seq_len
tmp_test_constructor.n_samples = tmp.n_samples
if random_seed != 0:
tmp_test_constructor.no_save = True
tmp_test_constructor.setUp()
model_original = copy.deepcopy(tmp_test_constructor.simplified)
states = defaultdict(list)
def record_state_hook(s, model, self, *args, **kwargs):
if isinstance(self, lsrl.LinState):
node_idx = model.topological_sort().index(self)
s[node_idx].append(copy.deepcopy(self._state))
# add the hook that records outputs to all nodes
for node in model_original.graph().nodes():
node.hooks["before_forward"].append(lambda self, *args, **kwargs: record_state_hook(states, model_original, self, *args, **kwargs))
# PREPARE THE NOISY MODEL
model_noisy = copy.deepcopy(tmp_test_constructor.simplified)
noise_activations = list()
# go through the nodes and each node's Matrix type attribute and add a little gaussian noise to it
# remember that the Matrices are scipy sparse matrices
if noise_level != 0:
for node in model_noisy.graph().nodes():
for attr in node.__dict__:
if isinstance(getattr(node, attr), lsrl.Matrix) and attr in ["A", "B", "b", "bias"]:
new_data = getattr(node, attr).matrix.data + np.random.normal(0, noise_level, getattr(node, attr).matrix.data.shape)
new_matrix = copy.deepcopy(getattr(node, attr))
new_matrix.matrix.data = new_data
setattr(node, attr, new_matrix)
# add a hook to the LinStates setting the state to the unnoised one.
# That's because we want to evaluate the effect of noising without complicating
# the analysis with error accumulation through the states. If didn't do that,
# the error rate would be monotonically increasing in the prompt length
def fix_state_hook(model, self, *args, **kwargs):
if isinstance(self, lsrl.LinState):
node_idx = model.topological_sort().index(self)
original_state = states[node_idx].pop(0)
self._state = original_state
for node in model_noisy.graph().nodes():
if isinstance(node, lsrl.LinState):
node.hooks["before_forward"].append(lambda self, *args, **kwargs: fix_state_hook(model_noisy, self, *args, **kwargs))
# EVALUATE AT ALL INPUTS
results_original = []
results_noisy = []
for full_input in full_inputs:
states.clear()
noise_activations.clear()
results_original.append(model_original(full_input[None, :]).numpy()[:, -tmp_test_constructor.seq_len:].flatten())
results_noisy.append(model_noisy(full_input[None, :]).numpy()[:, -tmp_test_constructor.seq_len:].flatten())
inner_results[(conditional_mode, noise_level)] = np.mean(np.average(np.array(expected).round() == np.array(results_noisy).round()))
return inner_results
modes = ["multiplicative", "direct", "optimized", "step_optimized", "multiplicative_optimized"]
results = Parallel(n_jobs=16)(delayed(process_one_target)(i, noise, mode) for i, noise, mode in tqdm.tqdm(list(product(range(n_targets), [0, 1e-3, 1e-5, 1e-7, 1e-9], modes))))
for res in results:
for k, v in res.items():
overall_results[k].append(v)
df = pandas.DataFrame({"conditional_mode": m, "Noise level": n, "Wrong outputs (%)": 100*(1-e)} for (m, n), l in overall_results.items() for e in l)
df.to_csv("noise_robustness_tok2tok.csv")
In [ ]:
plt.figure()
df = pandas.read_csv("noise_robustness_tok2tok.csv")
sns.stripplot(
hue="conditional_mode",
y="Wrong outputs (%)",
x="Noise level",
dodge=True,
data=df,
hue_order=["direct", "optimized", "step_optimized", "multiplicative", "multiplicative_optimized"],
# log_scale=True,
alpha=0.5,
)
plt.grid(axis="y", which="both")
plt.savefig("noise_effect_tok2tok.pdf", bbox_inches="tight")
In [ ]: