Efficient Two-Stage Bit Rate Matching for JPEG AI via Directional Model Selection and Log-Domain Root Finding

JPEG AI, the first international standard for learned image compression, supports a continuous bit-rate range through four models whose operating points are adjusted by a Gain Unit. To reach a target bit rate, the Bit Rate Matching (BRM) process of the reference Verification Model evaluates all four models and then performs a bisection-based iterative search without early termination. In this work, a two-stage method called MB-BRM is proposed. It replaces the exhaustive model scan with a directional model selection strategy and performs the search for the selected model as a log-domain root-finding problem that reuses the results of the first stage, terminates early once the tolerance is satisfied, and falls back to bisection as a safeguard. The method is evaluated on JPEG AI VM 7.0 in four configurations and compared with the reference method and a previously proposed accelerated BRM method. MB-BRM reduces the average number of trials from 13.4 to 4.1, which is approximately 70% fewer than the reference method and 43% fewer than the accelerated method, and shortens the BRM time by a factor of 2.2–2.5 relative to the reference. It selects the same model as the reference in all cases and reaches the target tolerance on the final bitstream in 93.6–98.3% of the cases, compared with 42–43% for the reference. The Bjøntegaard delta rate analysis shows that these gains are obtained while preserving rate–distortion performance.