JoVox: Joint Hierarchical Flow Matching for Lightweight Voxel Generation

JoVox. Joint Hierarchical Flow Matching for Lightweight Voxel Generation. By Xiangjun Tang and two others, from King Abdullah University of Science and Technology. Published on the seventh of November.

Voxels are attractive partly because they are light. Generating them at high fidelity without losing that lightness is still an open problem in computer vision. Current generators land on one side of a trade: high fidelity but computationally heavy, or lightweight but lower in fidelity. Hierarchical methods do keep voxels light, and their weakness is specific. The coarse levels decide where every finer cell is allowed to appear, yet they are generated first, with the least information about the shape, and they are never revisited. The result is often plausible but wrong, and nothing downstream can undo it.

The authors' key insight runs against the usual direction of a hierarchy: let the finer cells take part in deciding the coarse ones. JoVox, a discrete flow matching framework, does this by learning a single joint distribution over all levels of the voxel hierarchy, instead of a separate distribution for each level. Under a joint distribution, a coarse cell can be predicted while finer cells nearby have already been generated, so the early, sparse decisions no longer have to be made in the dark.

Two further pieces sit inside that framework. The first is a two-level prediction. Rather than outputting cells one level finer at each step, the network outputs cells two levels finer at once. That halves the number of levels the network has to process, and the paper reports it yields voxels of higher fidelity as well.

The second is an unmasking strategy aimed at masked diffusion models. Once finer cells have been generated, it revisits the earlier decisions and corrects inconsistencies between them. Where the hierarchical methods described earlier fixed their coarse levels for good, this gives those levels a second look with more of the shape in view.

JoVox outperforms the baselines on both unconditional and image-conditioned voxel generation. The sharpest comparison is with TRELLIS.2, described as the current state of the art. Against it, JoVox reports eight times higher training throughput. A forward pass costs one fifteenth of the FLOPs. Training runs on four GPUs, where TRELLIS.2 uses thirty-two.

And on quality, rather than trading it away for those savings, JoVox comes out ahead, with a Chamfer distance eleven percent lower.

Both sides of the original trade move together here. The compute figures point toward the lightweight end, the Chamfer distance toward the high-fidelity end, and what connects them is a single change of view: treat the hierarchy as one joint distribution, so that what is learned at the fine levels can flow back up to the coarse ones.