How chihshengj/robust-steganography works
1. This repository is a Python package for two related steganography systems: an embeddings-based hidden-message channel and a watermarking-based channel. The codebase is organized around reusable core primitives, dataset/prompt generation, evaluation experiments, and saved measurement artifacts rather than a single application runtime.
2. Packaging is centered in pyproject.toml, requirements.txt, and uv.lock, with source code under src/. The repository also includes many checked-in experiment outputs and figure files, which makes it both a library and a research artifact store.
3. The embeddings system lives under src/embeddings/. Its architecture is split into:
- core/: encoder logic, error correction, hash functions, and task-specific “systems” such as story_system, summary_system, litreview_system, topicQA_system, unit_test_system, cot_system, and simulation helpers.
- config/: prompt templates and constants used to steer text generation for each covertext mode.
- utils/: embedding retrieval, paraphrasing, PCA helpers, sampling, and random-seed utilities.
This subsystem is designed to embed information into semantically meaningful text and recover it robustly after text transformations.
4. The embeddings pipeline has a strong PCA/training-data component under src/pca/. Different dataset domains have their own generation scripts, including enron_emails, litreview, summary, unit_test, cot_reasoning, and creative_stories. The Enron branch is the most developed pipeline: it parses emails, creates paraphrases, computes embeddings, combines chunks, and prepares training data for a PCA-based hash model. This is the main bridge between raw text corpora and robust semantic hashing.
5. The watermarking system lives under src/watermarks/ and is built as a composable stack:
- models/: language-model adapters such as gpt2.py and shakespeare_nanogpt.py
- prf/: pseudorandom functions like AES- and HMAC-based PRFs
- perturb/: distribution perturbation strategies (smooth_perturb, harsh_perturb)
- core/: the embedder and extractor that coordinate watermark insertion and recovery
- covertext/: covertext-length calculation and related support
- attacks/: text-modification attacks used to evaluate robustness
- utils/: config, helpers, and debugging support
The central boundary here is model/tokenizer + PRF + perturbation feeding the embedder/extractor pair.
6. Watermark generation and recovery are separated cleanly from the model implementation. The model layer supplies token distributions; the PRF maps keys and history to pseudo-random token preferences; perturbation alters those distributions; the embedder samples watermarked text; and the extractor replays the same key/history logic to recover counters or bits. That separation is important for diagramming because each layer can be swapped independently.
7. The attack layer is first-class, not auxiliary. ngram_shuffle, synonym, translation, and paraphrase represent different threat models, from local edits to LLM-based paraphrasing. In the repository’s framing, local edit attacks are within the watermark robustness target, while GPT-style paraphrasing is explicitly stronger and is one of the key stress tests for the embeddings-based system.
8. experiments/ is the orchestration layer for running end-to-end studies. It is organized as a phased pipeline:
- phase1_generate.py for generating texts
- phase2_metrics/ for token counts, perplexity, and stegoanalysis
- phase3_attacks.py for adversarial transforms
Shared helpers in experiments/utils/ handle prompt expansion, system selection, IO, metrics, perplexity, and token counting. The dry_run/ subproject contains a concrete miniature dataset with prompts, generated texts, feature extraction, and summary results.
9. measurements/ is the analysis and reporting layer. It contains comparison scripts, watermark/embedding measurement scripts, preliminary studies, and recovery tests across several content domains. The figures/ tree stores output tables and logs, showing that this repo tracks experiment artifacts as part of the architecture, not just source code.
10. The repository has a strong test/example split. examples/ demonstrates both embeddings and watermark workflows, while src/watermarks/tests/ provides unit tests for attacks, covertext behavior, and PRF bias. Together, these indicate that the key quality gates are functional correctness of the primitives and empirical robustness of the end-to-end pipelines.
11. External dependencies matter mainly where they define system boundaries: OpenAI-style LLM calls appear in paraphrasing-based attacks and some generation/judging workflows; Semantic Scholar requests are used in the lit-review PCA dataset pipeline; and Hugging Face-style language models are implied by the GPT2/NanoGPT watermark adapters. Those integrations are not the core of the library, but they are part of the experimental data flow.
12. In diagram form, the repo naturally splits into three layers: source libraries (src/embeddings, src/watermarks, src/pca), experiment orchestration (experiments, measurements), and artifacts/results (experiments/.../results, measurements/figures). The main data flow is: dataset/prompt generation → text generation or embedding/watermark encoding → attack/transform steps → metric extraction/recovery evaluation → saved summaries and figures.
Main entry points: watermarks/core Watermark core · experiments Experiments · embeddings/core Embeddings core · pyproject.toml Packaging
Components 19
Source libraries
Reusable steganography and PCA code under src/.
- Embeddings core · Encoding, error correction, hash functions, and task-specific embedding systems.
src/embeddings/core - Embeddings config · Prompt templates and constants that steer covertext generation.
src/embeddings/config - Embeddings utils · Embedding lookup, paraphrasing, PCA helpers, sampling, and seed utilities.
src/embeddings/utils - PCA pipelines · Domain-specific dataset generation and PCA training preparation.
src/pca - Enron flow · Email parsing, paraphrasing, embedding chunking, and PCA dataset preparation.
src/pca/enron_emails - Other domains · Dataset generation for lit review, summary, unit test, CoT reasoning, and creative stories.
src/pca/litreview - Watermark models · Language-model adapters that supply token distributions.
src/watermarks/models - Watermark PRF · AES/HMAC pseudorandom functions that map key and history to token preferences.
src/watermarks/prf - Perturbation · Smooth and harsh perturbation strategies applied to model distributions.
src/watermarks/perturb - Watermark core · Embedder and extractor that coordinate watermark insertion and recovery.
src/watermarks/core - Covertext support · Covertext-length calculation and related support code.
src/watermarks/covertext - Attack layer · N-gram shuffle, synonym, translation, and paraphrase attacks.
src/watermarks/attacks
Experiments
End-to-end generation, metrics, attacks, and dry-run orchestration.
- Experiments · Phase-based end-to-end study runner for generation, metrics, and attacks.
experiments - Exp utils · Prompt expansion, system selection, IO, metrics, perplexity, and token counting.
experiments/utils - Dry run · Concrete miniature dataset, feature extraction, and summary outputs.
experiments/dry_run - Examples · Embeddings and watermark walkthrough scripts.
examples
Artifacts
Checked-in results, figures, and measurement outputs.
- Measurements · Comparison scripts, recovery tests, preliminary studies, and reporting code.
measurements - Figures · Saved tables, logs, and analysis outputs from experiments and measurements.
measurements/figures
Other components
- Packaging · Repository packaging and dependency specification.
pyproject.toml
Connections 21
- Packaging → Embeddings core: packages · Defines and distributes the embeddings library.
- Packaging → PCA pipelines: packages · Includes PCA and dataset preparation code.
- Packaging → Watermark core: packages · Includes the watermarking library.
- Embeddings config → Embeddings core: drives · Prompt templates steer task-specific generation.
- Embeddings utils → Embeddings core: supports · Provides retrieval, paraphrase, PCA, and sampling helpers.
- PCA pipelines → Enron flow: specializes · Enron is the main dataset-processing branch.
- PCA pipelines → Other domains: specializes · Other domain dataset generators feed the PCA pipeline.
- Enron flow → Embeddings utils: uses · Relies on embeddings and paraphrase utilities.
- Watermark models → Watermark core: feeds · Supplies token distributions to the embedder and extractor.
- Watermark PRF → Watermark core: keys · Maps key and history into pseudo-random token preferences.
- Perturbation → Watermark core: modulates · Alters distributions before sampling and recovery.
- Covertext support → Watermark core: constrains · Provides covertext-length logic used during embedding.
- Attack layer → Watermark core: tests · Creates modified texts for robustness evaluation.
- Exp utils → Experiments: supports · Shared orchestration helpers used by all phases.
- Experiments → Dry run: includes · Dry run is a concrete miniature study configuration.
- Experiments → Attack layer: runs · Phase 3 applies adversarial transforms.
- Experiments → Embeddings core: runs · Phase 1 generation and embedding workflows use the embeddings system.
- Experiments → Measurements: feeds · Outputs metrics and study results for analysis.
- Measurements → Figures: stores · Persistent tables, logs, and comparison outputs.
- Examples → Embeddings core: demonstrates · Shows embeddings workflows in runnable scripts.
- Examples → Watermark core: demonstrates · Shows watermark generation and attack examples.