Tropical tensor network for MPE

[GiggleLiu]

Two key component packages:

• omeco
• tropical-gemm

[Morningliumengkun]

Plan: Tropical Tensor Network for MPE (Issue #50)

This plan proposes a tropical tensor network approach for Most Probable Explanation (MPE). It includes an algorithm outline, key functions, and detailed API design. It assumes the two key components requested in issue #50 are:

• Tropical tensor primitives with einsum-style contraction (max-plus semiring ops, log-domain tensors).
• Tensor network contraction engine (factor graph → contraction order → binary contraction tree).

If the issue comment names different components, this plan can be updated accordingly.

---

1. Problem Statement

Goal: Compute MPE (MAP over all variables) for UAI-style factor graphs by converting factors to a tensor network and contracting in the tropical semiring (max, +) instead of the standard sum-product semiring.

Inputs:

• UAI factor graph (variables, cardinalities, factors).
• Optional evidence (fixed variable assignments).

Outputs:

• MPE assignment (variable → value).
• MPE score (log-probability / log-potential).

---

2. Algorithm (Simple Process)

1. Load model (UAI factors) and apply evidence.
2. Convert factors to log-tensors (safe log).
3. Build tensor network:
   • Nodes = factors
   • Indices = variables
4. Choose contraction order:
   • Greedy heuristic (min-fill / min-degree).
5. Build a binary contraction tree from the order.
6. Tropical contraction (einsum):
   • Combine two tensors at a time with a binary einsum.
   • Eliminate summed variables via max.
7. Track argmax to recover assignment:
   • Store backpointers at each binary contraction.
8. Assemble MPE assignment and final score.

---

3. Key Components (Two Packages)

Component A: Tropical Tensor Primitives (Einsum-Centric)

Purpose: define tensors over max-plus semiring and operations used in MPE contraction.

Core capabilities:

• tropical_einsum (binary contraction in max-plus semiring)
• safe_log (zero to -inf)
• argmax_trace (backpointers for assignment recovery)

Component B: Tensor Network Contraction Engine

Purpose: build the tensor network from factors and contract using a chosen order.

Core capabilities:

• build_network(factors)
• choose_order(graph, heuristic="min_fill")
• build_contraction_tree(order)
• contract_tree(tree, einsum_fn)

---

4. Proposed Project Structure (Aligned with Example)

tropical_tn/
├── README.md
├── requirements.txt
├── setup.py (optional)
├── src/
│   ├── __init__.py
│   ├── primitives.py         # tropical einsum, safe log, backpointer helpers
│   ├── network.py            # network building and data structures
│   ├── contraction.py        # order selection + binary contraction tree
│   ├── mpe.py                # top-level MPE APIs
│   └── utils.py              # shared helpers (ordering, shapes, checks)
├── tests/
│   ├── __init__.py
│   ├── test_primitives.py
│   ├── test_contraction.py
│   └── test_mpe.py
├── examples/
│   └── asia_network/
│       ├── main.py
│       └── model.uai
└── docs/
    ├── mathematical_description.md
    ├── api_reference.md
    └── usage_guide.md

---

5. Key Functions (High-Level)

• build_tropical_factors(model, evidence=None)
• choose_contraction_order(network, heuristic="min_fill")
• build_contraction_tree(order)
• tropical_einsum(a, b, index_map)
• contract_tree(tree, einsum_fn, track_argmax=True)
• recover_mpe_assignment(backpointers, tree)
• mpe_tropical(model, evidence=None, order=None)

---

5.1 Naming and Operator Conventions

• Follow NumPy/PyTorch naming patterns (einsum, tensor, contract, reduce).
• Prefer operator overloading for semiring composition:
  • Use a + b to mean tropical "sum" (i.e., max).
  • Use explicit tropical_einsum for contraction instead of separate mul/reduce APIs.
• Avoid introducing extra surface APIs when a standard name or operator conveys intent.

---

6. Detailed API (Proposed)

6.1 primitives.py

Function	Signature	Description
safe_log	safe_log(tensor: torch.Tensor) -> torch.Tensor	Convert to log domain; maps zeros to -inf.
tropical_einsum	tropical_einsum(a, b, index_map, track_argmax=True)	Binary einsum-style contraction in max-plus semiring; returns tensor and optional backpointer metadata.
argmax_trace	argmax_trace(meta, assignment)	Decode assignment using saved backpointer metadata.

6.2 network.py

API	Signature / Fields	Description
class TensorNode	vars: tuple[int, ...], values: torch.Tensor	Node representing a factor in log-domain.
build_network	build_network(factors: list[Factor]) -> list[TensorNode]	Convert factors into tensor nodes with variable scopes.

6.3 contraction.py

Function	Signature	Description
choose_order	choose_order(nodes: list[TensorNode], heuristic="min_fill") -> list[int]	Select elimination order over variable indices.
build_contraction_tree	build_contraction_tree(order)	Build a binary contraction tree from the order.
contract_tree	contract_tree(tree, einsum_fn, track_argmax=True)	Contract along the tree using the tropical einsum.

6.4 mpe.py

Function	Signature	Description
mpe_tropical	mpe_tropical(model: UAIModel, evidence: dict[int, int] | None = None, order=None)	Returns (assignment_dict, score, info); info includes order, num_nodes, num_elims.

---

7. Example (Demo)

Input: examples/asia_network/model.uai
Evidence: optional (none for baseline)
Output: MPE assignment and log-score

Example script tropical_tn/examples/asia_network/main.py:

• Load model
• Run mpe_tropical
• Print assignment and score

---

8. Risks and Mitigations

• Large tensor blow-up: use greedy order, limit scope size, optional mini-buckets.
• Argmax recovery bugs: add unit tests for backpointers on tiny models.
• Precision issues: use float64 and safe log to avoid -inf propagation errors.

---

9. Deliverables

• Tropical tensor network package (tropical_tn/).
• MPE demo script.
• Documentation snippet for usage and algorithm summary.

[GiggleLiu]

It is a very practical proposal in general, great job @Morningliumengkun . I have a few comments to improve this proposal:

1. I recommend implementing an einsum operation over two tensors, similar to: https://github.com/under-Peter/OMEinsum.jl , instead of exporting multiple APIs trace, reduce and mul. So we have: uai -> tropical tensor network -> optimize contraction order and create a binary contraction tree -> contract the network with the binary eninsum.
2. Please follow the numpy or pytorch naming principle, consider overloading the "+" instead of creating a new API.
3. The example is good, maybe we can use the famous asia network for demonstration purpose, show it in the documentation.

I think with these changes, this proposal is good to go.