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

Component Backbone Trainable part in this checkpoint
Compressor meta-llama/Llama-3.2-1B-Instruct full weights (compressor/)
Connector Linear(2048→4096) – ReLU – Linear(4096→4096) full weights (connector.pt)
Decoder mistralai/Mistral-7B-Instruct-v0.2 LoRA r=64 on all linear layers, plus resized embed_tokens and lm_head (decoder_state.pt)
  • 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:

Stage Rows Samples
1 (from scratch) [0, 500k) 0.5M
2 [500k, 2M) 1.5M
3 (this checkpoint) [2M, 4M) 2M

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)

Metric Value
Autoencoding reconstruction, Rouge-L / exact match 1.00 / 1.00
Autoencoding loss with <AE> 0.0002
Text-continuation loss 1.184
Mixed eval loss / Rouge-L ~0.95–1.02 / ~0.70

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.

Downloads last month
61
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for julia-bel/comp_mistral_pretrain_4M

Adapter
(664)
this model

Dataset used to train julia-bel/comp_mistral_pretrain_4M

Paper for julia-bel/comp_mistral_pretrain_4M