CET-ViT v4 — Hierarchical Vision Transformer with Sparse Token Grouping

Learning multi-scale visual hierarchies via sparse token assignment. Causal-emergence theory motivated this design; it is not validated by these experiments. See Corrections and Disclosures.

---

Overview

CET-ViT is a hierarchical Vision Transformer in which micro-tokens (image patches) are routed into a small set of macro-tokens through a learned sparse assignment (V-CEO module). Each token attaches to only its top-3 macro slots rather than diffusing across all 32. The macro-tokens are then processed by a second encoder and fused back with the micro-scale representation.

The demonstrated driver of the results below is sparse token grouping. An auxiliary spectral regularizer, originally motivated by causal-emergence theory, is included in the training objective but contributes substantially less (see Component attribution) and its theoretical interpretation did not survive verification.

What this repository provides

• Trained checkpoints for CET-ViT v4 and a no-EI ablation
• Full training logs (per-epoch, with timestamps)
• External baselines trained under an identical protocol (Swin, ToMe)
• Multi-seed results with variance on two datasets
• Analysis scripts, including the diagnostics that produced the corrections below

---

Corrections and Disclosures

An earlier version of this model card and the associated manuscript made two claims that we subsequently verified to be incorrect. Both are documented here in full rather than quietly removed. The forward pass and the released checkpoints are unchanged — these are errors of description and attribution, not of the trained weights.

1. The Dynamic-K estimator is non-functional

The architecture contains a DynamicKEstimator module intended to predict a per-image token budget K via Gumbel-softmax over candidates {4, 8, 16, 32}. This module never trains. Four independent checks confirm it:

Its output is not used anywhere in the reported results. Its parameters remain in the released checkpoints as vestigial dead weights.

Consequence for the "Dynamic K" finding. The per-image K values reported below (whale ≈ 9 vs. hamster ≈ 16) are real, but they are computed post-hoc from the sparse assignment matrix by counting slots whose normalised usage exceeds 5% at inference time. They are a property of the learned sparse assignment, not the output of a learned K-estimator. The observation stands; the mechanism attributed to it was wrong.

Reproduce with:

python -m experiments.analysis.probe_kgrad
python -m experiments.analysis.check_checkpoint_k \
    --ckpt checkpoints/cet_vit_v4_best_ep287_76.54.pth --data ./data

2. The reversibility metric is inverted relative to its description

The EI loss was described as a double-well: an entropy floor preventing rank-1 collapse, and a "reversibility ceiling" preventing over-diffusion. The metric reversibility = σ₁ / Σσᵢ is high when the assignment is concentrated and low when it is diffuse — the opposite of the stated reading:

Both loss terms therefore penalise concentration; neither penalises over-diffusion. There is no double-well and no implemented "[0.20, 0.45] emergence zone". Furthermore, at the model's actual operating point (reversibility ≈ 0.24) both ReLU terms evaluate to zero, which is consistent with the small logged ei value (≈0.017 at convergence) and implies the EI loss acts in practice mainly through its third term, the macro-distinctiveness (degeneracy) penalty.

Reproduce with:

python -m experiments.analysis.probe_reversibility

3. Claims withdrawn

• "First application of causal emergence theory to vision transformers." Withdrawn as an overclaim. More importantly, the mechanism — assigning tokens to a smaller set of learned groups, yielding emergent semantic regions without segmentation supervision — has close, well-established prior art in Slot Attention (Locatello et al., 2020) and GroupViT (Xu et al., 2022), the latter also using Gumbel-softmax assignment. See Related Work.
• "Maintaining meaningful causal structure (reversibility 0.241)." Withdrawn; the metric does not measure what this sentence claims.
• Hoel-style Effective Information is never computed. SVD reversibility is a spectral proxy, not determinism/degeneracy in Hoel's sense.

---

Results

CIFAR-100 (from scratch, 300 epochs, identical protocol)

The released checkpoint (76.54%) is a single favourable run; the 3-seed mean of 76.05 ± 0.09 is the number to compare against.

Tiny-ImageNet-200 (64×64, patch 4 → 256 tokens, 200 epochs)

Welch t-test: t = 6.99, p = 0.017; seed ranges do not overlap (CET 57.27–57.54 vs. Swin 56.37–56.58).

The advantage narrows at higher resolution

The margin over a size-matched hierarchical baseline shrinks by ~77% when moving to higher resolution and more classes. The Tiny-ImageNet win is statistically significant, but the benefit of sparse hierarchical grouping appears largest in the small-image regime. We report this explicitly as a limitation rather than a footnote.

Component attribution

The ablation table is the basis for centring this work on sparse grouping rather than on the emergence-inspired loss:

Most of the gain is attributable to sparse token assignment, not to the EI term.

Calibration

The model is systematically under-confident. A single global temperature fixes almost all of it, at no cost to accuracy:

Temperature fitted on one half of the validation set and evaluated on the other.

Reversibility vs. accuracy: no significant correlation

Per-class reversibility and per-class accuracy over 100 CIFAR-100 classes:

• Pearson r = 0.170, p = 0.092
• 95% bootstrap CI for r = [−0.01, 0.34] (crosses zero)
• Spearman ρ = 0.180, p = 0.065

We observe a weak, non-significant positive trend and do not claim it confirms any hypothesis. It is reported as exploratory.

Slot-usage K vs. object complexity

Computed post-hoc from the assignment matrix (see Disclosure 1) — not from a learned estimator:

OOD robustness (vs. no-EI ablation)

Single-run numbers; noise robustness is worse, not better.

Baseline not included

LTM-Transformer. No official implementation is publicly available. Our own reimplementation did not train stably within our compute budget (two attempts, ~100 GPU-hours; the second collapsed to chance accuracy with every optimizer step skipped for non-finite gradients). We report ToMe as the representative token-merging baseline and omit LTM rather than publish a number we cannot verify as faithful to the original method.

---

Architecture

Input Image (32×32)
  │
  PatchEmbed (2×2 patches → 256 tokens)
  │
  MicroEncoder (4-stage Swin-like, 29.8M params)
  │
  ┌──▼─────────────────────────────────────────┐
  │              V-CEO Module                  │
  │  Sparse top-3 softmax assignment           │  ← the operative mechanism
  │  → S ∈ ℝ^{B×N×K_max}   (K_max = 32)        │
  │  → h_macro = Sᵀ · h_micro                  │
  │  [DynamicKEstimator: present but inert —   │
  │   receives no gradient, see Disclosure 1]  │
  └──┬─────────────────────────────────────────┘
  │
  MacroEncoder (2.7M params)
  │
  CrossScaleAttention (0.3M params)
  │
  DeepEncoder + Fusion Head
  │
  Classification (100 classes)

The V-CEO assignment adds only 0.026M parameters (0.1%).

EI loss, as actually implemented

# Term 1: entropy floor
norm_entropy    = H(σ) / log(K)
entropy_penalty = ReLU(0.5 - norm_entropy)      # fires when CONCENTRATED

# Term 2: labelled "reversibility ceiling"
reversibility   = σ₁ / Σσᵢ                       # HIGH = concentrated
rev_penalty     = ReLU(reversibility - 0.45)     # also fires when CONCENTRATED

# Term 3: macro distinctiveness (degeneracy)
degen_loss      = mean off-diagonal cosine similarity of h_macro

L_EI = entropy_penalty + 2.0 * rev_penalty + 0.1 * degen_loss

Note both ReLU terms respond to concentration, and at the model's operating point (reversibility ≈ 0.24) both are zero. See Disclosure 2.

---

Model Card

---

Repository Structure

cet-vit-v4-cifar100/
├── checkpoints/          # LFS — run `git lfs pull` (≈380 MB each)
│   ├── cet_vit_v4_best_ep287_76.54.pth
│   └── ablation_no_ei_best_ep271_74.15.pth
├── figures/
├── logs/                 # per-epoch JSONL with timestamps
├── results/
└── src/
    ├── config.py
    ├── engine.py
    ├── data/             # cifar100.py (loaders + mixup_batch)
    ├── utils.py          # optimizer / scheduler / checkpoint helpers
    ├── models/           # cet_vit, vceo, dynamic_k, encoders
    ├── losses/           # svd_ei_loss, total_loss, pred_loss
    ├── probing/          # analysis scripts
    └── scripts/          # train_cifar_v4.py, train_ablation_no_ei.py

---

Quick Start

Checkpoints are stored with Git LFS:

git lfs install
git clone https://huggingface.co/Dhruv1000/cet-vit-v4-cifar100
cd cet-vit-v4-cifar100
git lfs pull
pip install -r requirements.txt

Load the model

import torch
from src.models.cet_vit import CETViT
from src.config import make_model_config

cfg = make_model_config(
    scale="base", img_size=32, patch_size=2, num_classes=100,
    drop_path_rate=0.3, drop_rate=0.1, attn_drop_rate=0.1,
    entropy_reg_weight=0.01, k_candidates=[4, 8, 16, 32],
)

model = CETViT(cfg)
ckpt = torch.load("checkpoints/cet_vit_v4_best_ep287_76.54.pth",
                  map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model"])
model.eval()

x = torch.randn(1, 3, 32, 32)
with torch.no_grad():
    logits, aux = model(x)

print(f"Predicted class: {logits.argmax().item()}")
print(f"S matrix shape : {aux['S'].shape}")   # [B, N_tokens, K_max]

# Number of active macro-regions = slot-usage count from S.
# Do NOT use aux['k_expected'] — it comes from the inert DynamicKEstimator
# and is not a meaningful per-image quantity (see Disclosure 1).
usage = aux['S'].sum(dim=1)
usage = usage / usage.sum(-1, keepdim=True).clamp(min=1e-8)
print(f"Active macro-regions K: {(usage > 0.05).sum(-1).item()}")

Train from scratch

python src/scripts/train_cifar_v4.py
python src/scripts/train_cifar_v4.py --debug    # 2-epoch smoke test

Reproduce the diagnostics

python -m experiments.analysis.probe_kgrad            # dead K-estimator
python -m experiments.analysis.probe_reversibility    # inverted metric
python -m experiments.analysis.check_checkpoint_k \
    --ckpt checkpoints/cet_vit_v4_best_ep287_76.54.pth --data ./data
python -m experiments.analysis.recompute_stats --data results/rev_acc_raw.json
python -m experiments.analysis.calibration --logits results/val_logits.pt

---

Reproducibility notes

Issues found in earlier revisions of this repository, now fixed. If you cloned before these fixes, re-pull:

• src/data/ and src/utils.py were missing from the published tree; the code could not train or run probing from a fresh clone.
• einops was imported by patch_embed.py but absent from requirements.txt.
• Several scripts in src/probing/ contained hard-coded absolute paths (/workspace/outputs, /workspace/data/cifar100) from the original training environment. These must be changed to local paths, or made configurable.
• Checkpoints are Git LFS objects: a plain git clone yields 134-byte pointer files. Run git lfs pull.

---

Related Work

The V-CEO assignment mechanism is closely related to existing object-centric and grouping approaches, which we did not adequately cite in earlier versions:

• Slot Attention — Locatello et al., NeurIPS 2020. Iterative competitive attention binding inputs to a small set of slots; produces object-centric groupings without segmentation supervision.
• GroupViT — Xu et al., CVPR 2022. Assigns segment tokens to fewer group tokens via Gumbel-softmax; the authors describe their grouping block as behaving like a single iteration of Slot Attention.

How V-CEO differs: a single feedforward top-3 sparse assignment, trained end-to-end from classification labels alone, with no iterative refinement (unlike Slot Attention) and no large-scale image–text contrastive pretraining (unlike GroupViT), at a cost of 0.026M parameters. We regard this as an efficiency/simplicity contribution rather than a novel grouping principle.

Other references:

• Hoel et al. (2013) — Quantifying causal emergence (motivation only; Hoel-EI is not computed in this work)
• Bolya et al. (2023) — Token Merging (ToMe), ICLR
• Liu et al. (2021) — Swin Transformer
• Marin et al. (2023) — Token Pooling in ViTs, WACV

---

Citation

@misc{cetvit2026,
  title  = {CET-ViT: Hierarchical Vision Transformer with Sparse Token Grouping},
  author = {Das, Dhruv Jyoti},
  year   = {2026},
  url    = {https://huggingface.co/Dhruv1000/cet-vit-v4-cifar100}
}

---

Original training: AMD MI300X · Reproduction and baselines: 4× NVIDIA A16 · PyTorch