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.

License: MIT

CIFAR-100

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:

Check	Result
Gradient flow	All 6 parameters return `grad = None` under the training loss
Trained weights vs. fresh init	Statistically indistinguishable (μ, σ match within noise) after 287 epochs
Training log	Logged `k_expected` sits at ≈15.0 (the near-uniform mean of {4,8,16,32}) for the entire run
Eval-mode output spread	Present, but reproduced at equal or greater magnitude by two independently-initialised untrained modules — an artifact of hard-argmax over a random projection, not evidence of learning

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:

bash

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:

Assignment matrix S	Reversibility	`ReLU(rev − 0.45)` fires?
Concentrated (near rank-1)	0.99	Yes
Diffuse (tokens across distinct slots)	0.04	No
Realistic top-3 sparse	0.05	No

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)

Model	Params	Val OA	Seeds
CET-ViT v4 (ours)	33.12M	76.05 ± 0.09	3
CET-ViT v4 (released checkpoint, ep287)	33.12M	76.54	1
Ablation (no EI)	33.12M	74.15	1
v3 (entropy-only EI, dense softmax)	33.12M	73.94	1
Swin (size-matched)	32.90M	72.14	1
ToMe (size-matched ViT, r=8)	33.43M	63.01	1

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)

Model	Params	Val OA	Seeds
CET-ViT v4 (ours)	33.2M	57.39 ± 0.11	3
Swin (size-matched)	32.98M	56.47 ± 0.11	2

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

Dataset	CET-ViT	Swin	Gap
CIFAR-100 (32×32, 100 cls)	76.05	72.14	+3.91
Tiny-ImageNet (64×64, 200 cls)	57.39	56.47	+0.92

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:

Configuration	Δ OA vs. baseline
Two-term EI, dense softmax	+0.67
Two-term EI, sparse (top-3) softmax	+1.72

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:

	ECE
Uncalibrated	0.1559
Temperature-scaled (T = 0.743)	0.0205

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:

Complexity	Example classes	Mean slot-usage K
Simple / uniform	whale, cloud, telephone	9.3–9.7
Average	most classes	12.86
Complex / detailed	hamster, tulip, tiger, baby	15.6–16.1

OOD robustness (vs. no-EI ablation)

Corruption	CET-ViT v4	No-EI	Δ
Gaussian Noise	8.27%	11.12%	−2.85
Blur	31.05%	28.68%	+2.37
Brightness ↓	75.75%	74.23%	+1.52
Contrast ↓	75.72%	74.34%	+1.38
Mean (corrupted)	47.70%	46.09%	+1.61

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

python

# 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

Property	Value
Architecture	CET-ViT Base
Parameters	33.12M
FLOPs	1.81G
Input	32×32×3 (CIFAR-100), 64×64×3 (Tiny-ImageNet, patch 4)
Best Val OA (released ckpt)	76.54% (epoch 287)
Val OA (3 seeds)	76.05 ± 0.09%
Best TTA OA (5-aug)	76.88%
ECE / after temperature scaling	0.1559 / 0.0205 (T = 0.743)
Original training hardware	AMD MI300X, ~129 s/epoch, ~10.7 h for 300 epochs
Reproduction hardware	4× NVIDIA A16 (15 GB, 62 W each), ~265–493 s/epoch
Framework	PyTorch ≥ 2.1

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:

bash

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

python

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

bash

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

Reproduce the diagnostics

bash

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

bibtex

@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