comp_mistral_pretrain_4M

comp_mistral_pretrain_4M

Built with Llama.

A compress-then-generate model in the PISCO style: a Llama-3.2-1B compressor turns a 128-token passage into 8 memory embeddings (compression rate 16), and a LoRA-adapted Mistral-7B-Instruct decoder generates from those embeddings. This checkpoint is the plain pretraining stage (autoencoding + text continuation) after 4M training samples; it has not been fine-tuned on any downstream task.

Architecture

• Compression rate 16, compressor window 128 tokens, decoder window 512 tokens.
• Two added special tokens: <MEM> (memory slots) and <AE> (autoencoding prompt); the decoder embedding matrix has 32,002 rows.
• Base decoder weights are not included: they are downloaded from mistralai/Mistral-7B-Instruct-v0.2 at load time.

Training

Three consecutive stages on EleutherAI/SmolLM2-135M-10B, each on fresh rows, 4M samples in total:

Objective: 50% autoencoding (reconstruct the passage from its memory embeddings, prompted with <AE>) and 50% text continuation. Stage 3 hyperparameters (training_config.yaml): lr 1e-4 with linear decay, warmup ratio 0.05, weight decay 0.1, effective batch 1,024 (4 GPUs × 4 × 64 gradient accumulation), 1 epoch = 1,953 steps, bf16, max grad norm 1.0, seed 42.

Results (held-out split of the pretraining data)

Reconstruction was already saturated after 2M samples; the last 2M samples mainly lower the continuation loss (1.207 → 1.184).

Usage

Loading requires the PISCO training code (pisco package) and access to the gated mistralai/Mistral-7B-Instruct-v0.2 repository. The loader reads a local directory, so download the snapshot first:

from huggingface_hub import snapshot_download
from pisco.model import PISCO

path = snapshot_download("julia-bel/comp_mistral_pretrain_4M")
model = PISCO.from_pretrained(path, load_decoder=True)

Files

config.json              PISCO configuration
compressor/              Llama-3.2-1B-Instruct compressor (safetensors, bf16)
connector.pt             connector weights
decoder_state.pt         Mistral-7B LoRA + resized embed_tokens / lm_head
compressor_tokenizer/    compressor tokenizer with <MEM>/<AE>
decoder_tokenizer/       decoder tokenizer with <MEM>/<AE>
training_config.yaml     stage-3 training configuration
LICENSE_LLAMA_3.2.txt    Llama 3.2 Community License (covers the compressor weights)

License

The compressor is a derivative of Llama 3.2 and is distributed under the Llama 3.2 Community License. The decoder adapter is derived from Mistral-7B-Instruct-v0.2 (Apache 2.0). Use of the base decoder is subject to its own terms.