Gravity Labeler: Architecture-Agnostic Data Profiling via Field Physics

Improve any model's training by 1.0–1.9% — including Transformers — by assigning smarter sequence lengths to your training data.

---

What is Gravity Labeler?

Gravity Labeler is a data profiling tool that analyzes your training corpus and assigns each document an information locality label:

Label	Meaning	Recommended seq_len
LOCAL	Spiky density — information is concentrated in short bursts	512
SENTENCE	Moderate density — information is sentence-scoped	1024
PARAGRAPH	Smooth density — information spans paragraphs	2048
DOCUMENT	Very smooth — information spans the full document	4096

These labels are derived from field physics (density coefficient of variation and field variance ratio), not document length. This means:

• A 10,000-token FAQ may only need local context → assigned LOCAL
• A 500-token legal clause may need full context → assigned DOCUMENT
• 43.8% of documents get different assignments than length-based methods

Works with any model — Transformer, Mamba, Gravity, or any other architecture.

---

Install

git clone https://github.com/chiangjw90/gravity-labeler.git
cd gravity-labeler
pip install -e .

Note: This package depends on gravity-nn. Install it first:

git clone https://github.com/chiangjw90/gravity-nn.git
cd gravity-nn && pip install -e . && cd ..

---

Quick Start

from gravity_labeler import GravityProbe, train_labeler, label_documents

# documents: list of lists of token ids
# e.g. [[101, 2054, 2003, ...], [101, 1996, ...], ...]

# 1. Train a small probe on your corpus (self-supervised, ~500 steps)
probe = GravityProbe(vocab_size=50257)
probe = train_labeler(probe, documents, vocab_size=50257, device="cuda")

# 2. Label all documents
all_stats = label_documents(probe, documents, device="cuda")

# 3. Inspect results
for doc in all_stats[:5]:
    print(doc["label"])       # LOCAL / SENTENCE / PARAGRAPH / DOCUMENT
    print(doc["rec_seq"])     # recommended sequence length: 512 / 1024 / 2048 / 4096
    print(doc["density_cv"])  # information concentration score

---

High-Level API: DataProfiler

For end-to-end profiling in one call:

from gravity_labeler import DataProfiler

profiler = DataProfiler(vocab_size=50257)
labels = profiler.profile(documents)
curriculum = profiler.build_curriculum(labels, token_budget=400_000_000)

print(curriculum["allocations"])
# {512: {'n_docs': 12000, 'token_budget': 50M, ...},
#  1024: {'n_docs': 35000, 'token_budget': 180M, ...}, ...}

---

Curriculum Training

After labeling, use the labels to build sequence-length buckets and train with curriculum ordering (short → long):

from gravity_labeler.curriculum import build_gravity_buckets, run_curriculum

# Build buckets from labels
buckets = build_gravity_buckets(
    documents, all_stats,
    bucket_sizes=[512, 1024, 2048, 4096]
)

# Train with curriculum
test_ppl, time_min, multi_results = run_curriculum(
    tag="my_experiment",
    make_model_fn=lambda V, N: my_model(V, N),
    buckets=buckets,
    lr=3e-4,
    device="cuda",
)

See examples/labeler_demo.py for a complete working example.

---

How It Works

A 5M-parameter GravityProbe is trained on your corpus for ~500 steps via next-token prediction. After training, it analyzes each document and extracts two physics-derived metrics:

• density_cv: Coefficient of variation of the information density field. High = spiky (local structure). Low = smooth (long-range structure).
• phi_var_ratio: Ratio of long-scale to short-scale field variance. High = long-range dependencies. Low = short-range only.

These two metrics form a 2D classification that assigns each document to a sequence length bucket. Compared to length-based assignment, this method correctly captures documents where length and information structure disagree.

---

Project Structure

gravity-labeler/
├── gravity_labeler/
│   ├── __init__.py       # Unified exports
│   ├── labeler.py        # GravityProbe model + train + label functions
│   ├── profiler.py       # DataProfiler: high-level end-to-end API
│   └── curriculum.py     # Bucket building + curriculum training loops
├── examples/
│   └── labeler_demo.py   # Complete end-to-end demo script
├── notebooks/            # Reproduction notebooks (coming in v0.2.0)
├── tests/
│   └── test_labeler.py   # Smoke tests
├── pyproject.toml
├── CITATION.cff
├── NOTICE
├── PATENT.md
└── LICENSE

---

Tests

pytest tests/ -v

---

Citation

If you use Gravity Labeler, please cite both papers:

Paper 2 (this work):

@article{chiang2026labeler,
  title={Gravity as Data Intelligence: Physics-Derived Information Locality Analysis for LLM Pretraining Data Routing},
  author={Chiang, Chia-Wei},
  year={2026},
  publisher={Zenodo},
  doi={10.5281/zenodo.20328022},
  url={https://doi.org/10.5281/zenodo.20328022}
}

Paper 1 (underlying Gravity architecture):

@article{chiang2026gravity,
  title={Gravity: A Physics-Inspired O(N) Framework with O(1) Streaming State Across Dimensions},
  author={Chiang, Chia-Wei},
  year={2026},
  publisher={Zenodo},
  doi={10.5281/zenodo.20273259},
  url={https://doi.org/10.5281/zenodo.20273259}
}

---

License

This software is offered under two licensing options:

1. AGPL-3.0 — for open-source use that complies with all AGPL-3.0 obligations, including source disclosure for network services (Section 13). Academic and research use is freely permitted under this option.

2. Commercial License — for use cases where AGPL-3.0 obligations are not acceptable (proprietary products, closed-source SaaS, embedded/hardware implementations). Contact chiangjw90@gmail.com

Patent pending: U.S. Provisional Patent Application filed March 2026. See PATENT.md for details.

---

Gravity Labeler: Because document structure is a physics problem, not a length problem.