# What Is Kompress-base and How Is It Trained? Inside Headroom's AI Compression Engine

> Discover Kompress-base, Headroom's AI text compression model trained on agentic traces for reversible token reduction. Learn how this ModernBERT innovation powers AI coding agents.

- Repository: [Tejas Chopra/headroom](https://github.com/chopratejas/headroom)
- Tags: deep-dive
- Published: 2026-06-12

---

**Kompress-base is a dual-head ModernBERT-based text compression model fine-tuned on agentic traces to power Headroom's reversible token reduction for AI coding agents.**

Kompress-base serves as the neural compression backend for the *Kompress* transform in the Headroom open-source repository. Unlike generic compression algorithms, this domain-specific model understands the structure of AI-generated artifacts—tool outputs, logs, and RAG chunks—to intelligently discard filler tokens while preserving semantic utility. The following sections dissect the architecture, training methodology, and runtime integration of Kompress-base as implemented in [`headroom/transforms/kompress_compressor.py`](https://github.com/chopratejas/headroom/blob/main/headroom/transforms/kompress_compressor.py).

## Core Architecture of Kompress-base

Kompress-base extends the ModernBERT transformer with a task-specific dual-head design optimized for binary token classification. The architecture balances fine-grained token decisions with contextual span awareness.

### ModernBERT Foundation

The encoder backbone is initialized from `answerdotai/ModernBERT-base`, a modernized BERT architecture optimized for efficiency. This base model provides **768-dimensional hidden representations** for each input token, serving as the feature extractor for the downstream compression heads.

### Token-Level Classification Head

The first prediction layer is a binary classifier implemented as `nn.Linear(hidden_size, 2)` in PyTorch. This **token-head** predicts a keep-versus-discard probability for every individual token in the input sequence. During inference, an argmax operation converts these logits into a hard binary mask indicating which tokens survive compression.

### Span-Level Importance Scoring

Complementing the token-head, a lightweight **1-D CNN** (`nn.Conv1d`) functions as the span-head. This convolutional layer produces per-token importance scores that capture contextual dependencies across contiguous regions of text. The span-head identifies semantically significant blocks—such as error messages or file paths—that might contain individual tokens classified as borderline by the token-head alone.

### Fusion Logic for Final Decisions

The final compression mask combines both head outputs through a confidence-boosting mechanism. The token-head's argmax provides the initial keep/discard decision. Tokens with probabilities between **0.3 and 0.5**—considered borderline—are promoted to "keep" status if the span-head assigns high importance scores to their surrounding context. This union operation ensures that critical contextual information is not fragmented during compression.

## Training Data and Methodology

Kompress-base is trained exclusively on **agentic traces**: real-world textual artifacts generated by AI coding agents including Claude, Codex, and Cursor deployments.

### The Agentic Traces Dataset

The training corpus comprises tool call outputs, system logs, RAG retrieval chunks, and intermediate agent thoughts collected from production Headroom environments. These traces represent the long, noisy content that typically overwhelms context windows during autonomous coding tasks.

### Ground Truth Generation via Heuristics and Annotation

Each training example pairs a raw trace with a **ground-truth compression** labeling essential versus expendable tokens. The ground truth derives from two sources:

- **Human annotation** of important versus filler content by domain experts.
- **Automatic heuristics** implemented in `_kompress_content_signature` that label tokens as essential when they contain file paths, error messages, identifiers, or high-utility code patterns.

This hybrid approach ensures the model learns both explicit structural cues and subtle semantic importance signals.

## Fine-Tuning Pipeline

The model undergoes supervised fine-tuning through a multi-stage pipeline optimized for the token classification objective.

### Preprocessing and Token Alignment

Raw texts are tokenized using the ModernBERT tokenizer, then split into word-aligned token groups. Each group inherits per-word keep flags derived from the ground-truth compressions, ensuring that subword tokens maintain consistent labels with their parent words.

### Loss Functions and Optimization

The training objective combines two loss components:

- **Binary cross-entropy loss** on the token-head outputs, optimizing the keep/discard decision accuracy.
- **Regression loss** on the span-head scores, teaching the CNN to produce calibrated importance weights for contiguous text regions.

Optimization uses the **AdamW** optimizer with gradient clipping and a learning-rate schedule featuring linear warm-up followed by cosine decay.

### Evaluation Metrics and Model Export

Validation on a held-out test set measures **token-level F1** and **must-keep recall** (the probability that ground-truth essential tokens are retained). The final checkpoint achieves **F1 ≈ 0.913** and **must-keep recall ≈ 0.977**, as reported in the ONNX artifacts comment block.

The trained PyTorch checkpoint is exported to multiple formats for flexible deployment:

- **Safetensors**: `model.safetensors` for standard HuggingFace loading.
- **ONNX**: `kompress-int8-wo.onnx` (INT8 quantized) and `kompress-fp32.onnx` (full precision) for optimized CPU inference.

These artifacts are hosted on HuggingFace under `chopratejas/kompress-v2-base`.

## Deployment in Headroom

The `KompressCompressor` class integrates the model into Headroom's transformation pipeline, handling lazy loading and inference optimization.

### Loading and Backend Selection

At runtime, the compressor downloads model weights and the ModernBERT tokenizer via `hf_hub_download_local_first`, caching artifacts locally to avoid repeated downloads. The implementation supports dual backends:

- **ONNX Runtime**: CPU-only execution using the quantized INT8 or FP32 models (controlled via `HEADROOM_KOMPRESS_BACKEND` and `HEADROOM_KOMPRESS_ONNX_FILENAME` environment variables).
- **PyTorch**: GPU-accelerated inference for high-throughput scenarios.

### Chunking and Inference

To manage long context windows, the compressor splits incoming text into **350-word chunks** (configurable via `KompressConfig.chunk_words`). Each chunk undergoes tokenization, dual-head inference, and mask assembly. The final compressed text concatenates only the tokens flagged for retention by the fusion logic.

The following pseudocode illustrates the core inference loop:

```python
from headroom.transforms.kompress_compressor import KompressCompressor

# Initialize with default or custom config

compressor = KompressCompressor()
compressor.load_model()  # Downloads from chopratejas/kompress-v2-base if needed

# Compress a long agent trace

compressed_text = compressor.compress(
    long_tool_output,
    chunk_words=350
)

# Result contains only semantically critical tokens

```

## Summary

- **Kompress-base** is a dual-head ModernBERT model (`answerdotai/ModernBERT-base`) fine-tuned specifically for compressing agentic traces.
- The architecture combines a **token-head** (binary classifier) and **span-head** (1-D CNN) to balance fine-grained and contextual importance scoring.
- Training leverages **agentic traces** with ground truth generated through human annotation and heuristics (`_kompress_content_signature`).
- The model achieves **F1 ≈ 0.913** and **must-keep recall ≈ 0.977** on the test split.
- Deployed via [`headroom/transforms/kompress_compressor.py`](https://github.com/chopratejas/headroom/blob/main/headroom/transforms/kompress_compressor.py), supporting both ONNX and PyTorch backends with configurable chunking at **350 words**.

## Frequently Asked Questions

### What makes Kompress-base different from standard text compression?

Unlike gzip or general-purpose language models, Kompress-base is domain-specific to AI agent outputs. It understands the semantic structure of tool calls, logs, and code snippets, allowing it to discard filler tokens while preserving paths, errors, and identifiers that downstream LLMs require for accurate reasoning.

### How does the dual-head architecture improve compression accuracy?

The **token-head** excels at local decisions but may miss contextual importance. The **span-head** (1-D CNN) provides surrounding context, rescuing borderline tokens (probability 0.3–0.5) that sit within important semantic spans. This fusion mechanism prevents the fragmentation of critical multi-token entities like file paths or stack traces.

### What model formats are available for production deployment?

Kompress-base is distributed as both **Safetensors** (`model.safetensors`) for PyTorch inference and **ONNX** (`kompress-int8-wo.onnx`, `kompress-fp32.onnx`) for optimized CPU-only deployment. The ONNX variants support INT8 quantization for reduced memory footprint without significant accuracy loss.

### How does Kompress-base handle long inputs exceeding context limits?

The implementation automatically fragments text into **350-word chunks** (configurable via `KompressConfig.chunk_words`), processes each chunk independently through the dual-head architecture, and reassembles the compressed segments. This chunking strategy ensures linear scaling with sequence length while maintaining local coherence.