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

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.

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:

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, 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}