diff --git a/configs/main_exp/mistral/self_gen_lv1_closed_qa_1_l2l.yaml b/configs/main_exp/mistral/self_gen_lv1_closed_qa_1_l2l.yaml
new file mode 100644
index 0000000..a941048
--- /dev/null
+++ b/configs/main_exp/mistral/self_gen_lv1_closed_qa_1_l2l.yaml
@@ -0,0 +1,31 @@
+# LoRA
+lora_r: 8
+lora_dropout: 0.0
+target_modules:
+  - down_proj
+
+use_kl_loss: true
+
+ctx_encoder_type: per_layer_activations
+n_latent_queries: 8
+num_blocks: 9
+num_self_attn_per_block: 0
+
+gradient_accumulation_steps: 11
+max_packed_inp_len: 6144
+max_packed_ctx_len: 6144
+
+# data
+train_ds_names:
+  - self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/fw_qa_v2/min_0_to_2000/train/*level_1*.parquet
+  - self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_0.0/pwc_compact
+  - self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/squad_compact
+  - self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/ropes_compact
+  - self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/drop_compact
+
+val_ds_names:
+  - squad
+  - pwc
+  - drop
+  - ropes
+  - self_gen/mistralai/Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet
diff --git a/scripts/main_exp/1-train.sh b/scripts/main_exp/1-train.sh
index 8c58a18..cc136af 100755
--- a/scripts/main_exp/1-train.sh
+++ b/scripts/main_exp/1-train.sh
@@ -2,23 +2,9 @@
 
 port=29051
 
-default_args=(
-    "--model_name_or_path=google/gemma-2-2b-it"
-    "--max_steps=50_000"
-    "--target_modules=down_proj"
-    "--lora_r=8"
-    "--eval_strategy=no"
-    "--max_qas_len=2048"
-    "--max_qas_per_sample=1"
-    "--per_rank_gen=True"
-    "--per_layer_processing=True"
-    "--gen_lora_l1_reg_coef=0.1"
-)
-
 uv run accelerate launch --config_file accelerate_config.yaml --main_process_port $port \
 --num_processes=8 --gpu_ids all train.py \
 configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml \
-"${default_args[@]}" \
 --model_name_or_path=google/gemma-2-2b-it \
 --target_modules=down_proj --lora_r=8 \
 --eval_strategy=no --max_qas_len=2048 --max_qas_per_sample=1 \
diff --git a/scripts/main_exp/2-train-chunk.sh b/scripts/main_exp/2-train-chunk.sh
new file mode 100755
index 0000000..4131e1a
--- /dev/null
+++ b/scripts/main_exp/2-train-chunk.sh
@@ -0,0 +1,30 @@
+#!/bin/bash
+
+port=29051
+
+uv run accelerate launch --config_file accelerate_config.yaml --main_process_port $port \
+--num_processes=8 --gpu_ids all train.py \
+configs/main_exp/self_gen_lv1_closed_qa_1_l2l.yaml \
+--model_name_or_path=google/gemma-2-2b-it \
+--target_modules=down_proj \
+--lora_r=8 \
+--eval_strategy=no \
+--max_qas_len=512 \
+--max_qas_per_sample=1 \
+--per_rank_gen=True \
+--per_layer_processing=True \
+--gen_lora_l1_reg_coef=0.1 \
+--max_steps=20000 \
+--gradient_accumulation_steps=16 \
+--max_packed_inp_len=1024 \
+--max_packed_ctx_len=2048 \
+--use_per_ctx_average_loss=True \
+--use_kl_loss=True \
+--quantize_ctx_encoder=True \
+--torch_empty_cache_steps=10 \
+--from_pretrained_checkpoint=train_outputs/runs/$RUN_NAME/checkpoint-80000/pytorch_model.bin \
+--max_ctx_chunk_len=512 \
+--min_ctx_chunk_len=25 \
+--num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' \
+--warmup_steps=2000 \
+--learning_rate=2e-5
diff --git a/scripts/main_exp/train-mistral.sh b/scripts/main_exp/train-mistral.sh
new file mode 100755
index 0000000..742a18b
--- /dev/null
+++ b/scripts/main_exp/train-mistral.sh
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+port=29051
+
+uv run accelerate launch --config_file accelerate_config.yaml --main_process_port $port \
+--num_processes=8 --gpu_ids all train.py \
+configs/main_exp/mistral/self_gen_lv1_closed_qa_1_l2l.yaml \
+--model_name_or_path=mistralai/Mistral-7B-Instruct-v0.2 \
+--target_modules=down_proj --lora_r=8 \
+--eval_strategy=no --max_qas_len=2048 --max_qas_per_sample=1 \
+--per_rank_gen=True --per_layer_processing=True --gen_lora_l1_reg_coef=0.1 \
+--max_steps=80000 --gradient_accumulation_steps=8 --max_packed_inp_len=4096 \
+--max_packed_ctx_len=4096 --use_per_ctx_average_loss=True --use_kl_loss=True \
+--quantize_ctx_encoder=True