Context Merging: from Tokens to Entities and Concepts

• prepare_dataset.py builds a local dataset of grouped embeddings from a base Qwen3 with a custom layer 0 that performs token grouping.
• train_custom_qwen3.py fine-tunes a customized Qwen3 that adds a small MLP adapter for grouped inputs, while freezing all weights except layer 0.
• inference_qwen3_merged.py runs end-to-end inference by first grouping with the base model, then generating with the trained model that understands grouped inputs. Includes perf metrics and estimated attention-memory savings.

How it works

1. Layer-0 grouping at prefill A custom decoder layer 0 computes attention on the full token sequence, clusters adjacent tokens using lightweight heuristics plus attention relations, then averages token vectors per group. The grouped result is added back to a residual projection and saved as grouped_hidden_states.
2. Dataset building The dataset builder swaps in the custom layer 0, feeds formatted prompts, extracts the stored grouped_hidden_states, and serializes them together with target responses.
3. Model training The training model wraps Qwen3 with a GroupedInputMLPAdapter that processes the grouped embeddings during prefill. Only layer 0 and the adapter are trainable; embeddings, upper layers, final norm, and LM head are frozen. Prefill uses grouped_inputs as inputs_embeds, then generation proceeds with past-key-values.
4. Inference The inference runner loads two models: a grouping model with the custom layer 0, and your trained model. It reports token compression, timing, and memory usage. Savings are also estimated with a simple attention-cost proxy that scales with sequence length squared.

Requirements

• Python packages: torch, transformers, datasets, tqdm, psutil. These are imported directly in the scripts.
• GPU is optional. Scripts detect CUDA and set dtype accordingly.

Repository layout

• prepare_dataset.py - dataset builder using custom layer 0 grouping.
• train_custom_qwen3.py - trainer for grouped-input Qwen3 with an MLP adapter, freezing all but layer 0.
• inference_qwen3_merged.py - two-stage inference runner with metrics.

1 Build the local dataset

• model_name="Qwen/Qwen3-0.6B"
• dataset_name="Magpie-Align/Magpie-Qwen2.5-Pro-1M-v0.1"
• output_dir="./processed_dataset"
• batch_size=1, max_samples=None, save_frequency=1000 Edit these in the constructor if you need to change them.

• processed_dataset.pkl - list of samples with inputs_embeds (grouped), response, and metadata.
• Additional metadata and sample previews are written alongside, for quick inspection.

2 Train the grouped-input model

• model_name="Qwen/Qwen3-0.6B"
• dataset_path="./processed_qwen3_dataset/processed_dataset.pkl"
• output_dir="./grouped_qwen3_checkpoint"
• batch_size=4, learning_rate=5e-4, num_epochs=3, warmup_steps=100
• Logging, eval, and checkpoint cadence are configurable.

• A GroupedInputMLPAdapter that takes grouped embeddings and returns adapted embeddings, normalized with RMSNorm.
• Only layer 0 and this adapter are trainable; everything else is frozen.

• Prefill: pass grouped_inputs via inputs_embeds with is_prefill=True.
• Then feed target response tokens while reusing past_key_values.

3 Run inference

Option A - standalone runner

• Loads a grouping model with the custom layer 0 and a trained inference model.
• Phase 1 groups tokens and reports compression. Phase 2 generates with the trained model.
• Reports compression ratio, memory reduction, total time, and tokens per second.

Option B - use the training script utilities

Parameters - quick reference

Dataset builder

• model_name - base HF model for grouping, default Qwen/Qwen3-0.6B.
• dataset_name - source HF dataset split, default Magpie-Align... Qwen2.5-Pro-1M.
• output_dir - where pickles and metadata go.
• max_samples - optional cap for quick tests.

Training

• dataset_path - path to processed_dataset.pkl.
• output_dir - where checkpoints are written.
• batch_size, learning_rate, num_epochs, warmup_steps - training hyperparams.
• Only layer 0 and the adapter are trainable. Verify with requires_grad settings in _freeze_layers.

Inference

• --checkpoint - path to trained checkpoint folder.
• --grouping_model - HF model name used for grouping.
• --instruction - user prompt, any language.
• --max_length, --temperature, --no_sample, --device.

Notes

• The custom layer 0 is installed by copying weights from the original layer 0, then replacing the module so it can compute groups and cache the grouped states.
• Grouping relies on simple rules over tokens like space and newline boundaries plus attention relations. You can tune the threshold in CustomQwen3Attention.

Troubleshooting

• CUDA memory spikes: reduce batch size during training or use fewer samples. Generation is incremental and reuses past-key-values.
• No grouped states found: ensure the custom layer 0 is used and is_initialized is reset before each prefill.
• Checkpoint not found: the inference loader expects pytorch_model.bin or model.safetensors in the checkpoint directory.

Why this can save memory

Citation

@misc{Kolomeitsev2025ContextMerging,
      title = {Context Merging: from Tokens to Entities and Concepts},
      author = {Konstantin Kolomeitsev},
      year = {2025}
}

Contact