1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
//! Training box embeddings on a 20+ entity taxonomy.
//!
//! This example trains axis-aligned box embeddings to represent hierarchical
//! relationships through geometric containment. A box A containing box B
//! means "B is-a A" (e.g., dog is-a mammal).
//!
//! Taxonomy (3 levels, 25 entities):
//! entity
//! animal
//! mammal: dog, cat, whale, bat
//! bird: eagle, sparrow, penguin
//! fish: salmon, tuna
//! plant
//! tree: oak, pine
//! flower: rose, tulip
//! vehicle
//! car, truck, bicycle
//!
//! The training uses direct coordinate updates with four passes per epoch:
//! 1. Positive containment: expand head boxes to contain their children (push min down, max up).
//! 2. Parent regularization: tighten each parent's boundaries toward its children's extent.
//! 3. Negative separation: for sibling and cross-branch pairs, shrink boxes to break
//! full containment on at least some dimensions.
//! 4. Leaf shrinkage: gently contract leaf boxes toward their center so they
//! develop varied, tighter volumes.
//!
//! This is a simplified approach; production systems would use backpropagation
//! through the containment probability.
//!
//! Reference: Vilnis et al. (2018), "Probabilistic Embedding of Knowledge
//! Graphs with Box Lattice Measures"
//!
//! Run: cargo run -p subsume --example box_training
use ndarray::Array1;
use std::collections::HashMap;
use subsume::ndarray_backend::NdarrayBox;
use subsume::HyperBox as BoxTrait;
fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("=== Box Embedding Training (25 entities, direct coordinate updates) ===\n");
// --- Define taxonomy as (head, tail) containment pairs ---
// head should contain tail (head is more general).
let containment_pairs: Vec<(&str, &str)> = vec![
// Level 0 -> 1
("entity", "animal"),
("entity", "plant"),
("entity", "vehicle"),
// Level 1 -> 2
("animal", "mammal"),
("animal", "bird"),
("animal", "fish"),
("plant", "tree"),
("plant", "flower"),
// Level 2 -> 3
("mammal", "dog"),
("mammal", "cat"),
("mammal", "whale"),
("mammal", "bat"),
("bird", "eagle"),
("bird", "sparrow"),
("bird", "penguin"),
("fish", "salmon"),
("fish", "tuna"),
("tree", "oak"),
("tree", "pine"),
("flower", "rose"),
("flower", "tulip"),
("vehicle", "car"),
("vehicle", "truck"),
("vehicle", "bicycle"),
];
// Collect all entity names
let mut entity_set = std::collections::HashSet::new();
for (h, t) in &containment_pairs {
entity_set.insert(*h);
entity_set.insert(*t);
}
let entity_names: Vec<&str> = {
let mut v: Vec<&str> = entity_set.iter().copied().collect();
v.sort();
v
};
let n_entities = entity_names.len();
println!("Entities: {}", n_entities);
println!("Containment pairs: {}\n", containment_pairs.len());
// --- Build taxonomy structure ---
let mut children_of: HashMap<&str, Vec<&str>> = HashMap::new();
let mut parent_of: HashMap<&str, &str> = HashMap::new();
for &(head, tail) in &containment_pairs {
children_of.entry(head).or_default().push(tail);
parent_of.insert(tail, head);
}
// Sibling index: position among siblings under the same parent.
let mut sibling_idx: HashMap<&str, usize> = HashMap::new();
for children in children_of.values() {
for (i, child) in children.iter().enumerate() {
sibling_idx.insert(child, i);
}
}
// Depth for each entity (root = 0).
let mut depth: HashMap<&str, usize> = HashMap::new();
for &name in &entity_names {
let mut d = 0;
let mut cur = name;
while let Some(&p) = parent_of.get(cur) {
d += 1;
cur = p;
}
depth.insert(name, d);
}
// --- Initialize box embeddings ---
//
// Hierarchy-aware initialization: assign each entity a position based
// on its branch in the taxonomy. Siblings get different offsets on a
// dedicated "separation dimension" so they start non-overlapping.
// Parents start large enough to cover their children's initial positions.
let dim = 8;
let mut boxes: HashMap<&str, (Array1<f32>, Array1<f32>)> = HashMap::new();
for &name in &entity_names {
let d = depth[name];
let half = match d {
0 => 5.0, // root: very large
1 => 3.0, // level 1: large
2 => 1.5, // level 2: medium
_ => 0.4, // leaves: small
};
// Build center: walk up the tree, accumulating sibling offsets.
// Each level uses a different dimension for separation.
let mut center = vec![0.0f32; dim];
let mut cur = name;
while let Some(&p) = parent_of.get(cur) {
let si = sibling_idx.get(cur).copied().unwrap_or(0);
let sep_dim = depth[cur] % dim; // which dim to separate on
center[sep_dim] += (si as f32) * 2.5;
cur = p;
}
let min_arr = Array1::from_vec(center.iter().map(|c| c - half).collect());
let max_arr = Array1::from_vec(center.iter().map(|c| c + half).collect());
boxes.insert(name, (min_arr, max_arr));
}
// --- Build negative (non-containment) pairs ---
//
// Sibling pairs: children of the same parent should not contain each other.
let mut negative_pairs: Vec<(&str, &str)> = Vec::new();
for children in children_of.values() {
for i in 0..children.len() {
for j in (i + 1)..children.len() {
negative_pairs.push((children[i], children[j]));
negative_pairs.push((children[j], children[i]));
}
}
}
let lr = 0.05;
let neg_lr = 0.04;
let shrink_lr = 0.002;
let parent_shrink_lr = 0.03;
let epochs = 300;
// --- Training loop ---
//
// Three passes per epoch:
// 1. Positive: expand head to contain tail (as before).
// 2. Negative: for sibling/cross-branch pairs, push boxes apart so
// A does not fully contain B.
// 3. Shrinkage: gently shrink leaf boxes toward their center so they
// don't all end up with identical volume.
println!(
"Training for {} epochs (dim={}, lr={}, neg_lr={})...\n",
epochs, dim, lr, neg_lr
);
// Identify leaf entities (those that never appear as a head).
let heads: std::collections::HashSet<&str> =
containment_pairs.iter().map(|(h, _)| *h).collect();
let leaves: Vec<&str> = entity_names
.iter()
.copied()
.filter(|n| !heads.contains(n))
.collect();
for epoch in 0..epochs {
let mut total_violation = 0.0f32;
// Pass 1: positive containment -- expand head to contain tail.
for &(head, tail) in &containment_pairs {
let (tail_min, tail_max) = boxes[tail].clone();
let (head_min, head_max) = boxes.get_mut(head).unwrap();
for d in 0..dim {
let margin = 0.05;
if head_min[d] > tail_min[d] - margin {
let violation = head_min[d] - (tail_min[d] - margin);
head_min[d] -= lr * violation;
total_violation += violation.abs();
}
if head_max[d] < tail_max[d] + margin {
let violation = (tail_max[d] + margin) - head_max[d];
head_max[d] += lr * violation;
total_violation += violation.abs();
}
}
}
// Pass 2: parent regularization -- tighten each parent's boundaries
// toward the actual extent of its children (plus margin). This
// prevents parents from growing far beyond what they need, which
// is the main cause of cross-branch contamination.
for (parent, children) in &children_of {
let mut child_min = vec![f32::MAX; dim];
let mut child_max = vec![f32::MIN; dim];
for &child in children {
let (cmin, cmax) = &boxes[child];
for d in 0..dim {
if cmin[d] < child_min[d] {
child_min[d] = cmin[d];
}
if cmax[d] > child_max[d] {
child_max[d] = cmax[d];
}
}
}
let margin = 0.1;
let (pmin, pmax) = boxes.get_mut(parent).unwrap();
for d in 0..dim {
let target_min = child_min[d] - margin;
let target_max = child_max[d] + margin;
if pmin[d] < target_min {
pmin[d] += parent_shrink_lr * (target_min - pmin[d]);
}
if pmax[d] > target_max {
pmax[d] -= parent_shrink_lr * (pmax[d] - target_max);
}
}
}
// Pass 3: negative separation -- for sibling pairs (a, b) where a
// should NOT contain b, push a's boundary inward on the single
// dimension where b is closest to escaping. We shrink a (not shift
// b) to avoid pushing b out of its actual parent.
for &(a_name, b_name) in &negative_pairs {
let (b_min_r, b_max_r) = boxes[b_name].clone();
// Find dimension where a covers b with smallest gap.
let (a_min_r, a_max_r) = &boxes[a_name];
let mut best_dim: Option<usize> = None;
let mut best_gap = f32::MAX;
for d in 0..dim {
if a_min_r[d] <= b_min_r[d] && a_max_r[d] >= b_max_r[d] {
let gap = (b_min_r[d] - a_min_r[d]).min(a_max_r[d] - b_max_r[d]);
if gap < best_gap {
best_gap = gap;
best_dim = Some(d);
}
}
}
if let Some(d) = best_dim {
let (a_min, a_max) = boxes.get_mut(a_name).unwrap();
let gap_min = b_min_r[d] - a_min[d];
let gap_max = a_max[d] - b_max_r[d];
// Push the closer boundary of a past b to break coverage.
if gap_min <= gap_max {
// Push a_min above b_min.
a_min[d] += neg_lr * (gap_min + 0.3);
} else {
// Push a_max below b_max.
a_max[d] -= neg_lr * (gap_max + 0.3);
}
total_violation += best_gap;
}
}
// Pass 4: shrink leaf boxes slightly toward their center.
// This prevents all leaves from having identical volume and
// encourages tighter, more specific representations.
for &leaf in &leaves {
let (leaf_min, leaf_max) = boxes.get_mut(leaf).unwrap();
for d in 0..dim {
let center = (leaf_min[d] + leaf_max[d]) * 0.5;
leaf_min[d] += shrink_lr * (center - leaf_min[d]);
leaf_max[d] -= shrink_lr * (leaf_max[d] - center);
}
}
// Re-enforce min < max invariant.
for (_name, (bmin, bmax)) in boxes.iter_mut() {
for d in 0..dim {
if bmin[d] >= bmax[d] {
let mid = (bmin[d] + bmax[d]) * 0.5;
bmin[d] = mid - 0.01;
bmax[d] = mid + 0.01;
}
}
}
if epoch % 50 == 0 || epoch == epochs - 1 {
println!(
" Epoch {:>4}: total_violation = {:.4}",
epoch, total_violation
);
}
}
// Build NdarrayBox instances for evaluation
let entity_boxes: HashMap<&str, NdarrayBox> = boxes
.iter()
.map(|(&name, (min, max))| {
let b = NdarrayBox::new(min.clone(), max.clone(), 1.0)
.expect("box construction should succeed after training");
(name, b)
})
.collect();
// --- Evaluate learned boxes ---
println!("\n--- Learned Box Volumes (larger = more general) ---\n");
let mut vol_pairs: Vec<(&str, f32)> = entity_boxes
.iter()
.map(|(&name, b)| (name, b.volume().unwrap_or(0.0)))
.collect();
vol_pairs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
for (name, vol) in &vol_pairs {
println!(" {:>12}: volume = {:.6e}", name, vol);
}
// --- Containment checks ---
println!("\n--- Containment Checks ---\n");
let checks: Vec<(&str, &str, &str, bool)> = vec![
// Positives (should have high containment probability)
("entity > animal", "entity", "animal", true),
("entity > vehicle", "entity", "vehicle", true),
("animal > mammal", "animal", "mammal", true),
("animal > bird", "animal", "bird", true),
("mammal > dog", "mammal", "dog", true),
("mammal > cat", "mammal", "cat", true),
("bird > eagle", "bird", "eagle", true),
("fish > salmon", "fish", "salmon", true),
("plant > tree", "plant", "tree", true),
("tree > oak", "tree", "oak", true),
("flower > rose", "flower", "rose", true),
("vehicle > car", "vehicle", "car", true),
// Negatives (should have low containment probability)
("dog > animal (reverse)", "dog", "animal", false),
("cat > dog (sibling)", "cat", "dog", false),
("animal > vehicle (cross)", "animal", "vehicle", false),
];
let mut correct = 0;
let total = checks.len();
for (label, head, tail, expect_high) in &checks {
let hb = &entity_boxes[head];
let tb = &entity_boxes[tail];
let p = hb.containment_prob(tb)?;
let ok = if *expect_high { p > 0.5 } else { p < 0.5 };
let status = if ok { "OK" } else { "FAIL" };
println!(" [{:>4}] {:<30} P = {:.3}", status, label, p);
if ok {
correct += 1;
}
}
println!(
"\nHierarchy accuracy: {}/{} ({:.0}%)",
correct,
total,
100.0 * correct as f32 / total as f32
);
println!("\nNotes:");
println!(" - This uses direct coordinate updates, not backpropagation");
println!(" - Negative separation pushes sibling/cross-branch boxes apart");
println!(" - Leaf shrinkage produces varied volumes (more specific = smaller)");
println!(" - Volume ordering (general > specific) emerges from containment constraints");
// See scripts/plot_box_concept.py for a visualization of box containment geometry.
Ok(())
}