Adaptive Discrete Communication Bottlenecks with Dynamic Vector Quantization

TL;DR

This paper introduces a method for adaptively adjusting the discretization level in vector quantization to improve multi-task learning and reinforcement learning performance by better matching data complexity.

Contribution

It proposes a novel approach to dynamically select discretization tightness in vector quantization based on input data, enhancing model flexibility and effectiveness.

Findings

Improved performance on visual reasoning tasks.

Enhanced reinforcement learning robustness.

Demonstrated benefits of adaptive discretization.

Abstract

Vector Quantization (VQ) is a method for discretizing latent representations and has become a major part of the deep learning toolkit. It has been theoretically and empirically shown that discretization of representations leads to improved generalization, including in reinforcement learning where discretization can be used to bottleneck multi-agent communication to promote agent specialization and robustness. The discretization tightness of most VQ-based methods is defined by the number of discrete codes in the representation vector and the codebook size, which are fixed as hyperparameters. In this work, we propose learning to dynamically select discretization tightness conditioned on inputs, based on the hypothesis that data naturally contains variations in complexity that call for different levels of representational coarseness. We show that dynamically varying tightness in communication bottlenecks can improve model performance on visual reasoning and reinforcement learning tasks.