Metadata-Version: 2.4
Name: ib-components
Version: 0.1.4
Summary: Independent building blocks for Information Bottleneck research: VIB compression, Riemannian optimization, and architectural improvements
Author: IB Components Team
License-Expression: MIT
Project-URL: Homepage, https://github.com/Fengrru/ib-components
Project-URL: Repository, https://github.com/Fengrru/ib-components
Project-URL: Issues, https://github.com/Fengrru/ib-components/issues
Project-URL: Documentation, https://github.com/Fengrru/ib-components#readme
Project-URL: Changelog, https://github.com/Fengrru/ib-components/blob/main/CHANGELOG.md
Keywords: information-bottleneck,vib,riemannian,deep-learning,neural-networks,pytorch
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Education
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Information Analysis
Classifier: Operating System :: OS Independent
Classifier: Typing :: Typed
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: pyyaml>=6.0
Requires-Dist: python-dotenv>=1.0.0
Provides-Extra: dev
Requires-Dist: pytest>=7.3.0; extra == "dev"
Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
Requires-Dist: pytest-mock>=3.12.0; extra == "dev"
Requires-Dist: black>=23.3.0; extra == "dev"
Requires-Dist: isort>=5.12.0; extra == "dev"
Requires-Dist: flake8>=6.0.0; extra == "dev"
Requires-Dist: mypy>=1.3.0; extra == "dev"
Requires-Dist: pre-commit>=3.3.0; extra == "dev"
Provides-Extra: ode
Requires-Dist: torchdiffeq>=0.2.3; extra == "ode"
Provides-Extra: transformers
Requires-Dist: transformers>=4.30.0; extra == "transformers"
Requires-Dist: einops>=0.7.0; extra == "transformers"
Provides-Extra: all
Requires-Dist: ib-components[ode,transformers]; extra == "all"
Provides-Extra: docs
Requires-Dist: sphinx>=7.0.0; extra == "docs"
Requires-Dist: sphinx-rtd-theme>=1.3.0; extra == "docs"
Requires-Dist: sphinx-autodoc-typehints>=1.24.0; extra == "docs"
Requires-Dist: sphinx-autobuild>=2024.0.0; extra == "docs"
Requires-Dist: myst-parser>=2.0.0; extra == "docs"
Dynamic: license-file

<div align="center">

# ib-components

**Modular building blocks for Information Bottleneck research in PyTorch**

[![CI](https://github.com/Fengrru/ib-components/actions/workflows/ci.yml/badge.svg)](https://github.com/Fengrru/ib-components/actions/workflows/ci.yml)
[![Python 3.9+](https://img.shields.io/badge/python-3.9+-3776AB?logo=python&logoColor=white)](https://www.python.org/)
[![PyTorch 2.0+](https://img.shields.io/badge/pytorch-2.0+-EE4C2C?logo=pytorch&logoColor=white)](https://pytorch.org/)
[![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)
[![Coverage](https://codecov.io/gh/Fengrru/ib-components/branch/main/graph/badge.svg)](https://codecov.io/gh/Fengrru/ib-components)

</div>

---

## Overview

`ib-components` is a modular library of Information Bottleneck building blocks for PyTorch. Unlike existing implementations ([VIB-pytorch](https://github.com/1Konny/VIB-pytorch), [information-bottleneck](https://github.com/xu-ji/information-bottleneck)) that bundle end-to-end training scripts tied to specific experiments, `ib-components` provides the **components themselves** as a library — compose them into your own architectures without forking.

### The Information Bottleneck Principle

The IB principle ([Tishby et al., 1999](https://arxiv.org/abs/physics/0004057)) formalizes the trade-off between compression and predictive power:

<p align="center">
  <img src="docs/_static/ib_principle.png" alt="IB Principle Formula" width="50%">
</p>

The **Variational Information Bottleneck** ([Alemi et al., 2017](https://arxiv.org/abs/1612.00410)) makes this tractable for deep networks:

<p align="center">
  <img src="docs/_static/vib_loss.png" alt="VIB Loss Formula" width="60%">
</p>

---

## Installation

```bash
# Core library
pip install ib-components

# With Neural ODE support (AttractorDynamics)
pip install "ib-components[ode]"

# With HuggingFace Transformers integration
pip install "ib-components[transformers]"

# All optional dependencies
pip install "ib-components[all]"

# Development install
git clone https://github.com/Fengrru/ib-components.git
cd ib-components
pip install -e ".[dev]"
```

**Requirements:** Python ≥ 3.9, PyTorch ≥ 2.0

---

## Quick Start

### VIB Compression (4 lines)

```python
import torch
from ib_components.compressor import LayerwiseVIBCompressor

compressor = LayerwiseVIBCompressor(hidden_size=768, latent_size=192, beta=0.001)

x = torch.randn(32, 128, 768)           # (batch, seq_len, hidden)
x_recon, info = compressor(x)            # reconstruct + latent info
losses = compressor.compute_ib_loss(info)  # IB loss

print(f"Total: {losses['loss']:.4f} | Recon: {losses['recon_loss']:.4f} | KL: {losses['kl_loss']:.4f}")
```

### Hierarchical Multi-Layer Compression

```python
from ib_components.compressor import HierarchicalVIBCompressor

compressor = HierarchicalVIBCompressor(num_layers=12, hidden_size=768, latent_ratio=0.25)

hidden_states = [torch.randn(32, 128, 768) for _ in range(12)]
reconstructed, infos = compressor(hidden_states)
total_loss = compressor.compute_total_loss(infos)
```

### Riemannian Optimization on Sphere Manifold

```python
from ib_components.compressor import project_to_tangent, sphere_exponential_map

point = torch.randn(1, 64)
point = point / torch.norm(point)  # normalize to unit sphere
grad = torch.randn(1, 64)

# Riemannian step — stays exactly on the sphere
riem_grad = project_to_tangent(grad, point)
riem_next = sphere_exponential_map(point, -0.1 * riem_grad)

# Euclidean step — drifts off the manifold
euc_next = point - 0.1 * grad

print(f"Riemannian norm: {torch.norm(riem_next):.6f}")  # 1.000000
print(f"Euclidean norm:  {torch.norm(euc_next):.6f}")    # ~1.07 (drifted)
```

### Adaptive Beta Scheduling

```python
from ib_components.improvements import AdaptiveBetaScheduler

scheduler = AdaptiveBetaScheduler(beta_min=0.0001, beta_max=0.01, schedule="sigmoid")

# High surprise → low beta (preserve information)
# Low surprise → high beta (aggressive compression)
for surprise in [0.0, 0.25, 0.5, 0.75, 1.0]:
    print(f"Surprise={surprise:.2f} → β={scheduler(surprise):.6f}")
```

---

## Features

### Components

| Module | Component | Description |
|--------|-----------|-------------|
| `compressor` | `LayerwiseVIBCompressor` | Single-layer VAE-style compressor with reparameterization |
| `compressor` | `HierarchicalVIBCompressor` | Multi-layer cascaded compressors with per-layer beta |
| `compressor` | `RiemannianOptimizer` | Sphere manifold optimization (tangent + exp map) |
| `compressor` | `SphereManifold` | Geodesic distance, parallel transport, normalization |
| `improvements` | `AdaptiveBetaScheduler` | Surprise-driven β adjustment (sigmoid/linear/exponential) |
| `improvements` | `KLAnnealingScheduler` | Linear/cosine/exponential KL warmup |
| `improvements` | `InformationGainDrive` | MINE-based mutual information estimation |
| `improvements` | `CrossLayerConsistency` | Frobenius norm alignment between adjacent layers |
| `improvements` | `AttractorDynamics` | Neural ODE with learnable attractor centers |

### Configuration

Type-safe dataclass configs with validation and YAML support:

```python
from ib_components.config import VIBConfig, ExperimentConfig

# Direct instantiation
config = VIBConfig(hidden_size=768, latent_size=192, beta=0.001)

# From YAML file
experiment = ExperimentConfig.from_yaml("config.yaml")

# To YAML
experiment.to_yaml("experiment_config.yaml")
```

---

## Benchmarks

*Synthetic data, `hidden_size=256`. Run with `python benchmarks/benchmark_compare.py`.*

| Metric | Value | Notes |
|--------|-------|-------|
| Compression ratio | **4×** (256→64) | 75% memory saved: 12 MB → 3 MB |
| Riemannian norm drift | **0.000001** | Euclidean step drifts to 1.258 (430,000× worse) |
| VIB KL loss | 139.9 | Compression vs reconstruction trade-off |
| Compressor parameters | 445K | Lightweight, ~0.13% of a 340M-param model |

<p align="center">
  <img src="docs/_static/riemannian_compare.png" alt="Riemannian vs Euclidean norm drift" width="90%">
  <br/><em>100 gradient steps on S<sup>255</sup>. Riemannian stays on the manifold; Euclidean drifts ~25% per step.</em>
</p>

<p align="center">
  <img src="docs/_static/compression_gain.png" alt="Compression ratio" width="90%">
  <br/><em>4× compression: 256-dim input → 64-dim latent, saving 75% memory.</em>
</p>

---

## Architecture

```
ib_components/
├── src/ib_components/
│   ├── compressor/              VIB encoder-decoder + Riemannian manifold
│   │   ├── vib_compressor.py    LayerwiseVIBCompressor, HierarchicalVIBCompressor
│   │   └── riemannian.py        RiemannianOptimizer, SphereManifold, tangent/expmap
│   ├── improvements/            Plug-and-play enhancements
│   │   ├── adaptive_beta.py     AdaptiveBetaScheduler, KLAnnealingScheduler
│   │   ├── info_gain.py         InformationGainDrive (MINE mutual information)
│   │   ├── cross_layer_consistency.py  CrossLayerConsistency (Frobenius alignment)
│   │   └── attractor_dynamics.py      AttractorDynamics (Neural ODE vector field)
│   ├── config.py                Type-safe dataclass configurations
│   ├── exceptions.py            Custom exception hierarchy
│   ├── logging.py               Structured logging (ComponentLogger)
│   └── utils/
│       ├── helpers.py           YAML config, logging, device, JSON, env, seeding
│       └── checkpointing.py     CheckpointManager, IncrementalCheckpointer
├── tests/                       Comprehensive test suite
├── benchmarks/                  Performance comparisons
├── examples/                    Runnable tutorials
└── docs/                        Sphinx documentation
```

Each component follows the same contract:

1. **Subclass `torch.nn.Module`** — drop-in PyTorch integration
2. **Accept `enabled: bool`** — set `False` for identity pass-through
3. **Return consistent outputs** — dict of loss terms or raw tensors

---

## Development

```bash
# Setup
git clone https://github.com/Fengrru/ib-components.git
cd ib-components
make dev

# Commands
make test       # Run tests (with coverage, ≥80% threshold)
make lint       # Run linting (black + isort + flake8)
make format     # Auto-format code
make typecheck  # Run mypy type checking
make check      # Run all checks (lint + typecheck + test)
make docs       # Build Sphinx documentation
make bench      # Run performance benchmarks
```

---

## References

| Paper | Year | Relevance |
|-------|------|-----------|
| [The Information Bottleneck Principle](https://arxiv.org/abs/physics/0004057) | 1999 | Foundational IB theory |
| [Deep Variational Information Bottleneck](https://arxiv.org/abs/1612.00410) | 2017 | VIB loss formulation |
| [MINE: Mutual Information Neural Estimation](https://arxiv.org/abs/1801.04062) | 2018 | Information gain estimation |
| [Neural Ordinary Differential Equations](https://arxiv.org/abs/1806.07366) | 2018 | Attractor dynamics |
| [Optimization on Matrix Manifolds](https://press.princeton.edu/books/hardcover/9780691132983/) | 2008 | Riemannian optimization theory |
| [KL Annealing for VAEs](https://arxiv.org/abs/1511.06349) | 2016 | KL warmup strategy |

---

## Citation

```bibtex
@software{ib_components2025,
  author       = {IB Components Team},
  title        = {ib\_components: Modular building blocks for Information Bottleneck research},
  year         = {2025},
  url          = {https://github.com/Fengrru/ib-components},
  version      = {0.1.3},
}
```

---

## Contributing

Contributions welcome! See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, coding standards, and the PR checklist.

## License

MIT — see [LICENSE](LICENSE) for details.
