PaperlayerView exact revision on arXiv

Source-linked AI summary

Tree-Structured Vector Quantization For Efficient And Progressive Image Compression

Xinkun Wang, Tianyi Xu, Qingyu Luo, Mingming Ma, Changzhe Jiao, Fu Li, Yi Niu

arXiv:2609.03641v1cs.CV

TL;DR

Most learned compression systems offer variable rates without ensuring that a lower-rate bitstream is a decodable prefix of a higher-rate one. Tree-VQ uses routed paths in a binary refinement tree, with progressive coding and prefix supervision. It produces smooth progressive rate–distortion behavior while retaining the efficiency benefits of structured tree routing.

  • Problem

    Variable-rate learned compression can produce different representations for different target rates without making the lower-rate representation a decodable prefix of the higher-rate one.

    Source: §1 · p. 1

  • Method

    Tree-VQ represents tokens as routed paths through a binary refinement tree and combines causal entropy modeling, rate-aware block refinement, and hierarchical prefix supervision.

    Source: §1 · pp. 2–3

  • Results

    Tree-VQ achieves smooth progressive rate–distortion behavior while preserving the efficiency benefits of structured tree routing.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Every prefix of a routed path is a valid reconstruction code, so additional branch symbols refine an existing representation rather than requiring re-encoding.

    Source: §1 · pp. 2–3

    The trained model has a relatively limited effective bitrate range, uses a single-scale latent design, and may miss globally better code assignments through greedy routing.

    Source: §5 · p. 9

Abstract

from arXiv · show

Vector-quantization based image compression has achieved strong rate--distortion performance, yet most of them still produce a separate compressed representation for each target bitrate. Such variable-rate behavior allows one model to operate at multiple rates, but it does not necessarily provide a progressive bitstream whose prefixes are themselves decodable and can be refined by appending additional bits. We propose \textbf{Tree-VQ}, a progressive tree-structured vector quantization framework for learned image compression. Tree-VQ organizes discrete codewords as a hierarchical binary tree and represents each latent token by a routed root-to-leaf path. Crucially, every prefix of this path corresponds to a valid quantized representation, so shallow internal nodes serve as coarse reconstruction codes and deeper nodes provide successive refinements. This allows a compressed image to be decoded from an early prefix and progressively improved as more branch symbols are received, rather than being re-encoded for different target rates. To make this structure practical for compression, we introduce a prefix-compatible tree entropy model that codes progressive continuation decisions and routed branch refinements using only causally available decoded contexts. We further use rate-aware refinement scheduling to decide which spatial blocks should receive additional tree bits under a given prefix budget, and hierarchical prefix supervision to ensure that internal nodes are directly decodable at low rates. Experiments show that Tree-VQ achieves a superior performance--efficiency trade-off, delivering the best perceptual compression results with much fewer parameters and lower latency than competing methods.

1 Introduction

Tree-VQ addresses the gap between model-level rate adaptability and bitstream-level progressivity by making each prefix of a routed tree path a valid reconstruction. Its progressive coding combines causally decodable entropy modeling, block-wise refinement, and supervision for useful low-rate representations.

1 Introduction

  • Most learned codecs provide model-level rate adaptivity, but different target rates need not be prefixes of one another that improve as bits arrive.Bitstream-level progressivity is relevant to network transmission, latency-sensitive decoding, scalable storage, and interactive visualization.Source: §1 · pp. 1–2
  • Tree-VQ lies near the lower-left of the DIV2K comparison with a small latency bubble, indicating a favorable trade-off among BD-LPIPS, parameter count, and run-time.Figure 1 reports relative BD-LPIPS against Tree-VQ over the overlapping bitrate range (bpp ≥0.25).Source: Figure 1 · p. 2, §1 · p. 2
  • Flat VQ indices are atomic symbols, so their prefixes usually do not act as meaningful coarse reconstruction codes.Tree-structured codebooks instead connect coarse internal prototypes to finer descendants and can reduce search complexity from O(N) to O(log N).Source: §1 · p. 2
  • Tree-VQ routes each latent token through a binary tree, making every root-to-path prefix a valid representation, from coarse shallow nodes to refined deeper nodes.Each additional branch decision provides a discrete refinement, and entropy-coded branch symbols form a prefix-decodable bitstream.Source: §1 · p. 2
  • Block-wise refinement layers use continuation symbols to select blocks for further refinement and branch symbols to extend token paths without replacing previously decoded paths.Later prefixes extend earlier paths rather than requiring re-encoding with a new depth map.Source: §1 · p. 2
  • The codec uses a prefix-compatible entropy model based on causally decoded contexts, rate-aware scheduling, and hierarchical prefix supervision for directly usable low-rate internal nodes.Scheduling prioritizes refinements with high distortion reduction per coding cost.Source: §1 · pp. 2–3

2 Related work

Flat VQ codebooks treat codewords as atomic prototypes, while progressive compression requires nested representations that remain decodable as bits arrive. Tree-VQ places this nesting in the codebook topology itself, unlike approaches that refine quantization grids, latent components, or enhancement streams.

2.1 Vector-quantized image representation and compression

  • Flat codebooks assign each latent vector to an unordered prototype, leaving related codewords without a shared coarse representation or a natural refinement mechanism.Increasing codebook size can improve representation capacity and reconstruction fidelity [48].Source: §2.1 · p. 3
  • VQ-VAE-2 [29] uses spatial hierarchies and RQ-VAE [17] additive residual sequences, whereas Tree-VQ uses prefixes of paths in a single codebook.The related approaches’ progressivity relies on adding vectors from separate codebooks or latent scales.Source: §2.1 · pp. 3–4
  • Tree-VQ makes every internal node on a token’s path a valid codeword, so each branch refines an existing representation rather than selecting an unrelated flat-codebook entry.This contrasts with flat VQ’s unordered prototypes, which do not naturally support progressive refinement.Source: §2.1 · pp. 3–4

2.2 Progressive learned image compression

  • Progressive codecs need nested bitstream prefixes that yield coarse reconstructions and refine the same representation, supporting preview, scalable storage, and adaptive transmission.The decoder may need to reconstruct an image before the full bitstream is available.Source: §2.2 · p. 4
  • PLONQ [23] nests scalar quantization grids, while scalable latent codecs use channels, residual masks, or enhancement streams; Tree-VQ instead nests the codebook itself.Each binary branch decision moves an already decoded code to one of its children.Source: §2.2 · p. 4

3 Method

Tree-VQ replaces an atomic flat codebook with a binary refinement tree whose path prefixes can be decoded as coarse-to-fine representations. Its progressive bitstream combines causal entropy coding, block-level refinement scheduling, and training and inference procedures designed to support useful intermediate reconstructions.

3.1 Prefix-decodable tree quantization

  • Each latent token follows a root-to-leaf path in a binary refinement tree, and every prefix provides a decodable coarse-to-fine representation.A balanced tree needs only L local binary routing decisions to reach among 2^L leaves, compared with a flat search over O(2^L) entries.Source: §3 · p. 4, §3.1 · p. 5

3.2 Embedded progressive bitstream and entropy modeling

  • The bitstream is arranged in successive refinement layers, where continuation symbols select blocks for further refinement and branch symbols extend their tokens’ paths.Layer-major raster ordering lets encoder and decoder use only previously decoded symbols and currently available prefixes.Source: §3.2 · p. 5
  • The prefix-compatible entropy model codes continuation and branch decisions from causal decoded context, without using future symbols, continuous latents, or a final target depth.Using the same distributions for training-time rate estimation and arithmetic coding aligns estimated prefix rate with realized bitrate.Source: §3.2 · pp. 5–6

3.3 Rate-aware progressive refinement scheduling

  • The scheduler ranks one-level block refinements by distortion benefit relative to coding cost, using encoder-side surrogates and entropy logits rather than repeated decoder passes.Its worst-case scheduling complexity is O(ML log M); the decoder only decodes transmitted symbols.Source: §3.3 · p. 6

3.4 Training objective and inference

  • Hierarchical prefix supervision decodes sampled tree depths and optimizes reconstruction quality at each, encouraging shallow nodes to preserve coarse content and deeper nodes to encode refinements.The codec’s progressive rate–distortion objective also trains performance at intermediate prefixes.Source: §3.4 · p. 7
  • Staged optimization first trains the encoder, decoder, and codebook with prefix supervision, then enables entropy modeling and scheduling for the full progressive objective.The paper says this procedure improves routing stability and prevents internal nodes from becoming unused routing artifacts.Source: §3.4 · p. 7
  • Encoder-side beam search selects token paths for the scheduler, while the decoder follows transmitted branches and reconstructs from the deepest decoded node at each prefix.Beam search changes neither the progressive bitstream format nor prefix decodability.Source: §3.4 · p. 7

4 Experiments

Across Kodak, CLIC2020, and DIV2K, Tree-VQ offers a smooth progressive rate–distortion frontier from one checkpoint and competitive results across operating points. Qualitative comparisons, routing benchmarks, and ablations examine reconstruction quality, efficiency, and component contributions.

4.1 Experimental setup

  • The evaluation trains on 300K randomly selected OpenImages images and compares Tree-VQ with learned codecs using distortion and perceptual metrics, with realized bitrate measured after arithmetic coding.Evaluation datasets are Kodak, CLIC2020, and DIV2K.Source: §4.1 · pp. 7–8

4.2 Main rate–distortion performance

  • From a single checkpoint, Tree-VQ traces a stable, smooth rate–distortion frontier: stronger rate penalties reduce bitrate, while weaker penalties select deeper representations and improve reconstruction quality.The operating points are obtained by truncating one embedded progressive bitstream at different prefix lengths.Source: §4.1 · p. 7, §4.2 · p. 8
  • Tree-VQ is competitive with strong learned codecs across Kodak, CLIC2020, and DIV2K, supporting a useful range of operating points without retraining for most bitrates.Its shared-tree, block-wise depth selection provides a different operating mechanism from separately trained rate points or global quantization control.Source: §4.2 · p. 8

4.3 Visual comparison

  • Tree-VQ produces visually competitive reconstructions, preserving major structures and object boundaries while maintaining reasonable texture fidelity at medium to high bitrates.The comparisons use Kodak images and local crops from a high-resolution DIV2K image, against M&S, HiFiC, MS-ILLM, and Control-GIC.Source: §4.3 · pp. 8–9

4.4 Model efficiency and scalability

  • At depth 16, Tree-VQ achieves up to 9195.1× speedup over flat vector quantization, with sub-millisecond latency in benchmarks reaching 65,536 target vectors.The latency comparison is reported across various tree depths and effective codebook sizes.Source: Table 1 · p. 9, §4.4 · p. 9

4.5 Ablation

  • Removing hierarchical prefix supervision causes the largest reconstruction-quality drop; removing the tree entropy model degrades rate performance, while Rebuild further improves the overall trade-off.The ablation uses the default adaptive-depth inference setting.Source: §4.5 · p. 9

5 Conclusion and future work

Tree-VQ uses a binary refinement tree, prefix-compatible entropy coding, rate-aware scheduling, and hierarchical prefix supervision to encode progressively decodable representations. The authors report smooth progressive rate–distortion behavior and efficient tree routing, while identifying bitrate range, latent scale, and greedy routing as limitations.

5 Conclusion and future work

  • Tree-VQ combines a binary refinement tree, prefix-compatible entropy model, rate-aware refinement scheduling, and hierarchical prefix supervision in a single embedded bitstream decodable at different prefix lengths.Different target bitrates do not require re-encoding.Source: §5 · p. 9
  • Experiments show smooth progressive rate–distortion behavior alongside the efficiency benefits of structured tree routing.Source: §5 · p. 9
  • The framework's effective bitrate range per trained model remains relatively limited, its single-scale latent design restricts multi-level representation capacity, and greedy routing may miss globally better code assignments.Proposed future work includes multi-scale quantization and stronger routing and refinement strategies.Source: §5 · p. 9

A Tree-VQ detailed implementation

Tree-VQ encodes each refinement layer without conditioning on a final truncation depth, using block continuation decisions and causal decoded contexts. During training, straight-through relaxation supports gradients while hard forward decisions preserve monotonic refinement.

A.1 Layer-wise prefix-compatible tree entropy model

  • Each layer is encoded without choosing a final truncation depth, so the first K layers can be decoded as a prefix before later layers arrive.For each block, a continuation symbol decides whether to refine; if it does, branch symbols for its tokens are encoded using causal contexts available to the decoder.Source: §A.1 · p. 13
  • The layer coding outcome is factorized using causal contexts computed from already decoded symbols, and the decoder knows the relevant context before decoding a branch symbol.Continuation value 1 refines the block to the next depth and triggers branch-symbol coding; value 0 leaves it at its current depth.Source: §A.1 · p. 13

A.2 Straight-through relaxation for progressive refinement decisions

  • Training uses a straight-through relaxation for layer-wise refinement decisions while keeping hard decisions in the forward pass.The relaxation uses a sigmoid and temperature parameter τ; progressive decisions remain monotonic, so blocks stay at their node or move deeper rather than replacing a decoded code.Source: §A.2 · pp. 13–14
  • A block's refinement score weighs estimated distortion reduction against incremental coding cost, and the hard forward pass refines it when the score meets the decision rule.The soft relaxation supports gradient propagation during training.Source: §A.2 · pp. 13–14

A.3 Full training objective and auxiliary regularization

The training objective combines progressive rate–distortion loss, hierarchical prefix supervision, and VQ stabilization. Auxiliary procedures supervise reconstructions across depths, encourage continuity along tree paths, and rebuild and monitor the tree during training.

A.3 Full training objective and auxiliary regularization

  • Tree-VQ's total training objective combines progressive rate–distortion loss, hierarchical prefix supervision, and VQ stabilization terms.The rate–distortion loss is computed over sampled prefix lengths, using reconstructions from the first K layers and the corresponding prefix coding cost.Source: §A.3 · p. 14
  • Hierarchical prefix supervision reconstructs images by truncating routed paths at multiple depths, encouraging shallow nodes to encode coarse structure and deeper nodes to provide refinement details.Source: §A.3 · p. 14
  • The multi-depth VQ loss compares encoder latents with centroids along their routed paths, updating centroids and applying commitment loss to encoder latents.Source: §A.3 · p. 14
  • Path consistency regularization penalizes large deviations between adjacent centroids to prevent abrupt parent–child jumps and stabilize progressive refinement.Source: §A.3 · p. 14
  • During the initial warm-up stage, top-down rebuilding updates centroids and recursively splits assigned latents to reduce severe early routing collapse.Rebuilding is used only during warm-up; subsequent centroid updates are gradient-based, and rebuilding adds no inference cost or bitstream-format change.Source: §A.3 · p. 15
  • Tree utilization is tracked with an effective utilization ratio and dead-node ratio to expose codebook imbalance and show how many branches are frequently selected.Utilization near 1 indicates balanced use of nodes at a depth; a small value indicates that only a small fraction of branches are frequently selected.Source: §A.3 · p. 15

B Experiment details

The experiment reports perceptual, distribution-based, no-reference, and auxiliary distortion measures of reconstruction quality.

B Experiment details

  • Reconstruction quality is assessed with full-reference perceptual scores LPIPS [46] and DISTS [7].The evaluation also includes no-reference NIQE and distribution-based FID [12] and KID, with PSNR as an auxiliary distortion metric.Source: B Experiment details · p. 15

B.1 Evaluation metrics

The evaluation uses several metrics because they capture different aspects of reconstruction quality.

B.1 Evaluation metrics

  • PSNR measures pixel-level fidelity, while LPIPS and DISTS better reflect perceptual similarity.The passage motivates reporting several measures rather than relying on a single score.Source: §B.1 · p. 15
  • NIQE evaluates no-reference naturalness, while FID and KID measure distribution-level realism on image patches.Source: §B.1 · p. 15
  • The paper reports all metrics jointly because different metrics emphasize different aspects of reconstruction quality.Source: §B.1 · p. 15

B.2 Evaluation of FID and KID.

FID and KID are computed on cropped patches for natural-image benchmarks, but not on Kodak because its patch count is considered insufficient for stable estimation.

B.2 Evaluation of FID and KID.

  • FID and KID follow a patch-based protocol used in generative image compression works [24] [27].Non-overlapping 256 × 256 patches are extracted, followed by additional overlapping patches shifted by 128 pixels in both spatial dimensions.Source: §B.2 · p. 16
  • The protocol yields 28,650 patches on CLIC2020 test and 6,573 patches on DIV2K validation.Source: §B.2 · p. 16
  • FID and KID are omitted on Kodak because its 24 images yield only 192 patches, insufficient for stable estimation.Source: §B.2 · p. 16

C Quantitative results

The section compares Tree-VQ with other methods using progressive rate–distortion curves on DIV2K; each Tree-VQ point comes from a prefix of the same embedded bitstream.

C Quantitative results

D.1 Uniform-depth tree reconstruction analysis

Uniform-depth reconstructions show that deeper Tree-VQ nodes refine coarse representations, while progressive examples demonstrate that later reconstructions can be decoded by extending the same bitstream. The checklist notes bitrate-range, architecture, routing, and code-release boundaries.

D.1 Uniform-depth tree reconstruction analysis

  • As depth increases from d = 2 to d = 10, reconstruction error decreases, especially near object boundaries and high-frequency regions.Intermediate nodes also serve as valid low-rate reconstruction codes. Figure 8 visualizes the error maps; Figure 7 compares distortion and statistical-fidelity metrics on CLIC2020.Source: §D.1 · p. 16, Figure 8 · p. 17, Figure 7 · p. 17

D.2 Progressive reconstructions

  • Later reconstructions on CLIC2020 and DIV2K come from appending branch-refinement symbols to the same bitstream, without re-encoding.The examples use the same progressive decoding prefixes.Source: §D.2 · p. 17, Figure 9 · p. 18, Figure 10 · p. 18

NeurIPS Paper Checklist

  • The checklist reports a limited bitrate range for one trained model, a single-scale latent design, possible suboptimality of greedy routing, and no released code repository in the current version.The paper uses public datasets.Source: NeurIPS Paper Checklist · pp. 19, 21
Loading 2609.03641v1…