Discrete Infomax Codes for Supervised Representation Learning

TL;DR

This paper introduces Discrete InfoMax Codes (DIMCO), a model that learns compact, discrete representations by maximizing mutual information with labels, reducing overfitting and improving efficiency in data processing.

Contribution

The paper proposes a novel probabilistic encoder for discrete codes that maximizes mutual information with labels and demonstrates its benefits for overfitting reduction and efficiency.

Findings

Shorter codes reduce overfitting in few-shot classification.

DIMCO codes are memory and retrieval efficient.

The model generalizes previous loss functions like cross-entropy.

Abstract

Learning compact discrete representations of data is a key task on its own or for facilitating subsequent processing of data. In this paper we present a model that produces Discrete InfoMax Codes (DIMCO); we learn a probabilistic encoder that yields k-way d-dimensional codes associated with input data. Our model's learning objective is to maximize the mutual information between codes and labels with a regularization, which enforces entries of a codeword to be as independent as possible. We show that the infomax principle also justifies previous loss functions (e.g., cross-entropy) as its special cases. Our analysis also shows that using shorter codes, as DIMCO does, reduces overfitting in the context of few-shot classification. Through experiments in various domains, we observe this implicit meta-regularization effect of DIMCO. Furthermore, we show that the codes learned by DIMCO are efficient in terms of both memory and retrieval time compared to previous methods.

Tables

Table 1. Table 1 : Top-1 accuracies under various compressed CIFAR-10/100 embeddings with size ( k , d 𝑘 𝑑 k,d ). The original embeddings were continuous 64 64 64 -dimensional embeddings ( k ≈ 2 32 , d = 64 formulae-sequence 𝑘 superscript 2 32 𝑑 64 k\approx 2^{32},d=64 ).

		CIFAR-10		CIFAR-100
		$d=4$	$d=16$	$d=4$	$d=16$
PQ	$k=2$	50.38	87.24	7.77	32.47
	$k=4$	86.86	90.92	18.78	52.02
	$k=8$	90.44	91.49	31.04	58.82
	$k=16$	91.02	91.52	43.16	61.00
DIMCO	$k=2$	64.43	88.85	10.03	33.46
(ours)	$k=4$	88.47	91.04	25.7	53.36
	$k=8$	91.29	91.45	46.22	58.84
	$k=16$	91.68	91.57	57.83	62.49

Table 2. Table 3 : Few-shot classification accuracies on the miniImageNet benchmark. † † \dagger denotes transductive methods, which are more expressive by taking unlabeled examples into account.

Method	$5$-way $1$-shot	$5$-way $5$-shot
† TPN	55.51 $\pm$ 0.86	69.86 $\pm$ 0.65
† FEAT	55.75 $\pm$ 0.20	72.17 $\pm$ 0.16
MetaLSTM	43.44 $\pm$ 0.77	60.60 $\pm$ 0.71
MatchingNet	43.56 $\pm$ 0.84	55.31 $\pm$ 0.73
ProtoNet	49.42 $\pm$ 0.78	68.20 $\pm$ 0.66
RelationNet	50.44 $\pm$ 0.82	65.32 $\pm$ 0.70
R2D2	51.2 $\pm$ 0.6	68.8 $\pm$ 0.1
MetaOptNet-SVM	52.87 $\pm$ 0.57	68.76 $\pm$ 0.48
DIMCO ($64,64$)	47.33 $\pm$ 0.46	61.59 $\pm$ 0.52
DIMCO ($64,128$)	53.29 $\pm$ 0.47	64.79 $\pm$ 0.57

Table 3. Table 4 : Image retrieval performance on CUB-200-2011 and Cars-196, measured by Recall@1. Memory is the number of bits that an embedding vector of each image uses. Time is seconds taken to retrieve a single query from database (5,924 and 8,131 images for CUB-200-2011 and Cars-196, respectively).

			CUB-200-2011		Cars-196
Method	($k,d$)	Memory [bits]	Recall@1	Time [s]	Recall@1	Time [s]
Binomial Deviance	(-, 128)	4096	57.25	16.37	72.53	21.86
Triplet	(-, 128)	4096	56.80	16.37	73.79	21.86
Proxy-NCA	(-, 128)	4096	56.19	16.37	75.94	21.86
DIMCO (ours)	(32, 32)	160	51.04	1.48	63.44	2.64
	(64, 64)	384	55.78	2.85	72.06	6.01
	(128, 128)	896	58.05	5.81	76.04	11.92
	(256, 256)	2048	58.90	12.20	77.32	16.04

Table 4. Table 5 : Top-1 accuracies under various compressed embedding sizes ( k , d 𝑘 𝑑 k,d ). The original embeddings were continuous 64 64 64 -dimensional embeddings, which corresponds to k = 2 32 , d = 64 formulae-sequence 𝑘 superscript 2 32 𝑑 64 k=2^{32},d=64 .

		CIFAR-10				CIFAR-100
		$d=2$	$d=4$	$d=8$	$d=16$	$d=2$	$d=4$	$d=8$	$d=16$
Product Quantization	$k=2$	21.88	50.38	78.68	87.24	2.84	7.77	15.24	32.47
	$k=4$	60.70	86.86	89.53	90.92	7.69	18.78	35.62	52.02
	$k=8$	90.46	90.44	90.98	91.49	15.28	31.04	50.74	58.82
	$k=16$	91.29	91.02	91.27	91.52	25.71	43.16	55.88	61.00
DIMCO (ours)	$k=2$	36.37	64.43	83.10	88.85	3.66	10.03	19.31	33.46
	$k=4$	62.92	88.47	90.78	91.04	11.83	25.7	37.4	53.36
	$k=8$	90.77	91.29	91.41	91.45	31.9	46.22	52.77	58.84
	$k=16$	91.49	91.68	91.46	91.57	46.17	57.83	61.11	62.49

Table 5. Table 6 : Top-1 accuracies of compressed embeddings under various sizes ( k , d 𝑘 𝑑 k,d ) on the ImageNet dataset. The original embeddings were continuous 2048 2048 2048 -dimensional embeddings, which corresponds to k = 2 32 , d = 2048 formulae-sequence 𝑘 superscript 2 32 𝑑 2048 k=2^{32},d=2048 .

	Product Quantization				DIMCO
	$d=2$	$d=4$	$d=8$	$d=16$	$d=2$	$d=4$	$d=8$	$d=16$
$k=2$	0.21	0.31	1.00	5.21	0.14	0.42	1.47	5.50
$k=4$	0.20	0.92	4.86	21.45	0.67	2.68	11.80	26.96
$k=8$	0.82	2.57	11.91	35.43	1.71	12.82	33.28	48.11
$k=16$	1.99	5.66	20.34	44.93	4.90	26.22	44.01	55.34

Table 6. Table 7 : Few-shot classification accuracies on the miniImageNet benchmark. Grouped according to backbone architecture. † † \dagger denotes transductive methods, which are more expressive by taking unlabeled examples into account.

Method	$5$-way $1$-shot	$5$-way $5$-shot
ConvNet (64-64-64-64)
† TPN (Liu et al., 2018)	55.51 $\pm$ 0.86	69.86 $\pm$ 0.65
† FEAT (Ye et al., 2018)	55.75 $\pm$ 0.20	72.17 $\pm$ 0.16
MetaLSTM (Ravi & Larochelle, 2016)	43.44 $\pm$ 0.77	60.60 $\pm$ 0.71
MatchingNet (Vinyals et al., 2016)	43.56 $\pm$ 0.84	55.31 $\pm$ 0.73
ProtoNet (Snell et al., 2017)	49.42 $\pm$ 0.78	68.20 $\pm$ 0.66
RelationNet (Sung et al., 2018)	50.44 $\pm$ 0.82	65.32 $\pm$ 0.70
R2D2 (Bertinetto et al., 2018)	51.2 $\pm$ 0.6	68.8 $\pm$ 0.1
MetaOptNet-SVM (Lee et al., 2019)	52.87 $\pm$ 0.57	68.76 $\pm$ 0.48
DIMCO ($64,64$)	47.33 $\pm$ 0.46	61.59 $\pm$ 0.52
DIMCO ($64,128$)	53.29 $\pm$ 0.47	64.79 $\pm$ 0.57
ResNet-12
† TPN (Liu et al., 2018)	59.46	75.65
† FEAT (Ye et al., 2018)	62.60 $\pm$ 0.20	78.06 $\pm$ 0.15
SNAIL (Mishra et al., 2017)	55.71 $\pm$ 0.99	68.88 $\pm$ 0.92
AdaResNet (Munkhdalai et al., 2017)	56.88 $\pm$ 0.62	71.94 $\pm$ 0.57
TADAM (Oreshkin et al., 2018)	58.50 $\pm$ 0.30	76.70 $\pm$ 0.30
MetaOptNet-SVM (Lee et al., 2019)	62.64 $\pm$ 0.61	78.63 $\pm$ 0.46
DIMCO ($128,64$)	54.57 $\pm$ 0.47	65.45 $\pm$ 0.31
DIMCO ($128,128$)	57.24 $\pm$ 0.44	69.31 $\pm$ 0.38

Equations

$$\left[\begin{array}{ccccc}{l}_{1,1}& l_{1,2}& l_{1,3}& \dots & l_{1,k}\\ l_{2,1}& l_{2,2}& l_{2,3}& \dots & l_{2,k}\\ ⋮& ⋮& ⋮& \ddots & ⋮\\ l_{d,1}& l_{d,2}& l_{d,3}& \dots & l_{d,k}\end{array}\right].$$

$$p_{i,j}=\frac{\exp (l_{i,j})}{{\sum }_{j=1}^k\exp (l_{i,j})}.$$

$$\displaystyle\mathbf{\widetilde{x}}=\Big{[}\overbrace{\operatornamewithlimits{\arg\max}_{j}p_{1j}}^{\widetilde{x}_{1}},\hskip 5.0pt\cdots,\hskip 5.0pt\overbrace{\operatornamewithlimits{\arg\max}_{j}p_{dj}}^{\widetilde{x}_{d}}\Big{]}^{\top}.$$

$$I(\tilde{X};Y)=H(\tilde{X})-H(\tilde{X}|Y).$$

$$H(p_1,\dots ,p_n)=\sum _{i=1}^n{p}_i\log p_i.$$

$$\hat{H}(\tilde{X})$$

$$\hat{H}(\tilde{X}|Y)=\sum _{y=1}^{c}p(Y=y)\hat{H}(\tilde{X}|Y=y),$$

$$\sum _{i=1}^d\log p_{i,{\tilde{x}}_i^{(s)}}({\mathbf{x}}^{(q)}).$$

$${\mathcal{L}}_{\text{ind}}=D_{KL}({\tilde{x}}_{i_1}\otimes {\tilde{x}}_{i_2}\Vert {\tilde{x}}_{i_1},{\tilde{x}}_{i_2}).$$

$$\mathcal{L}$$

$$=-H(\tilde{X})+H(\tilde{X}|Y)+\lambda {\mathcal{L}}_{\text{ind}}.$$

$$\underset{n\in \{1,\dots ,N\}}{\mathrm{argmax}}\left[\log p_{2,j_2}({\mathbf{x}}^n)+\log p_{4,j_4}({\mathbf{x}}^n)\right],$$

$$\mathrm{xent}(Y,X)=-{\mathbb{E}}_{y\sim Y,x\sim X}\left[\log q(\hat{Y}=y|\tilde{X}(x))\right],$$

$$\mathrm{xent}(Y,X)$$

$$\mathbb{E}\left[\log p(\tilde{\mathbf{x}}|y)-\log p(\tilde{\mathbf{x}})\right]$$

$$\mathcal{L}(\tau ,\theta )=-{\mathbb{E}}_{T\sim \tau }\left[I(\tilde{X}(X_T,\theta );Y_T)\right].$$

$$\hat{\mathcal{L}}(T^{1\text{mathchar}58n},\theta )=-\frac{1}{n}\sum _{i=1}^n\hat{I}(\tilde{X}(X_{T^i},\theta );Y_{T^i}).$$

$$\mathcal{L}(\tau ,\theta )-\hat{\mathcal{L}}(T^{1\text{mathchar}58n},\theta )\le O\left(\sqrt{\frac{d_{\Theta }}{n}\log \frac{n}{d_{\Theta }}}\right)+$$

$$O\left(\frac{|\tilde{X}|\log (m)}{\sqrt{m}}\right)+O\left(\frac{|\tilde{X}||Y|}{m}\right).$$

$$\mathcal{J}(p(\tilde{\mathbf{x}}|\mathbf{x}))=I(\tilde{X};X)-\beta I(\tilde{X};Y),$$

$$\mathrm{xent}(Y,\tilde{X})={\mathbb{E}}_{\mathbf{y}\sim Y,\tilde{\mathbf{x}}\sim \tilde{X}}\left[-\log q(\mathbf{y}|\tilde{\mathbf{x}},\varphi )\right].$$

$$\mathrm{argmin}\mathrm{xent}(Y,\tilde{X})$$

$$=$$

$${\mathcal{L}}_{\text{triplet}}={\Vert {\tilde{\mathbf{x}}}_q-{\tilde{\mathbf{x}}}_p\Vert }_2^2-{\Vert {\tilde{\mathbf{x}}}_q-{\tilde{\mathbf{x}}}_n\Vert }_2^2,$$

$$\mathbb{E}\left[-{\mathcal{L}}_{\text{triplet}}\right]$$

$$\approx$$

$$=$$

$$-\log \frac{\exp ({\tilde{\mathbf{x}}}^{\top }{\tilde{\mathbf{x}}}^{+})}{\exp ({\tilde{\mathbf{x}}}^{\top }{\tilde{\mathbf{x}}}^{+})+{\sum }_{i=1}^{N-1}\exp ({\tilde{\mathbf{x}}}^{\top }{\tilde{\mathbf{x}}}_i)}.$$

$$\mathbb{E}\left[\log p_p(\tilde{\mathbf{x}})-\log p(\tilde{\mathbf{x}})\right]=I(\tilde{X};Y)$$

$$\hat{\mathcal{L}}(T^{1\text{mathchar}58n},\theta )=-\frac{1}{n}\sum _{i=1}^n\hat{I}(\tilde{X}(X_{T^i},\theta );Y_{T^i}).$$

$$\mathcal{L}(\tau ,\theta )-\hat{\mathcal{L}}(T^{1\text{mathchar}58n},\theta )\le O\left(\sqrt{\frac{d_{\Theta }}{n}\log \frac{n}{d_{\Theta }}}\right)+O\left(\frac{|\tilde{X}|\log (m)}{\sqrt{m}}\right)+O\left(\frac{|\tilde{X}||Y|}{m}\right)$$

$$|I(\tilde{X};Y)-\hat{I}(\tilde{X};Y)|\le \frac{(3|\tilde{X}|+2)\log (m)\sqrt{\log (4/\delta )}}{\sqrt{2m}}+\frac{(|Y|+1)(|\tilde{X}|+1)-4}{m}$$

$$\left|I\left(\tilde{X}(X_T,\theta );Y_T\right)-\hat{I}\left(\tilde{X}(X_T,\theta );Y_T\right)\right|\le O\left(\frac{|\tilde{X}|\log (m)}{\sqrt{m}}\right)+O\left(\frac{|\tilde{X}||Y|}{m}\right).$$

Full text

Machine Learning, ICML

\DeclareTextFontCommand\bemph\bfseries

1 Introduction

Metric learning and few-shot classification are two problem setups that test a model’s ability to classify data from classes that were unseen during training. Such problems are also commonly interpreted as testing meta-learning ability, since the process of constructing a classifier with examples from new classes can be seen as learning. Many recent works (Hoffer & Ailon, 2015; Movshovitz-Attias et al., 2017; Snell et al., 2017; Oreshkin et al., 2018) tackle this problem by learning a continuous embedding ($\mathbf{\widetilde{x}}\in\mathbb{R}^{n}$) of datapoints. Such models compare pairs of embeddings using e.g., Euclidean distance to perform nearest neighbor classification. However, it remains unclear whether such models effectively utilize the entire space of $\mathbb{R}^{n}$.

Information theory provides a framework in which we can effectively ask such questions about representation schemes. In particular, the \bemphinformation bottleneck principle (Tishby et al., 2000; Shwartz-Ziv & Tishby, 2017) formalizes the optimality of a representation. This principle states that the optimal representation $\widetilde{X}$ is one that maximally compresses the input $X$ while also being predictive of labels $Y$. From this viewpoint, we see that the previous methods which map data to $\mathbb{R}^{n}$ focus on being predictive of labels $Y$ but not on compressing $X$.

The degree of compression of an embedding is the number of bits it reflects about the original data. Note that for continuous embeddings, each of the $n$ numbers in a $n$-dimensional embedding requires $32$ bits; it is unlikely that unconstrained optimization of such embeddings use all of these $32n$ bits effectively. We propose to resolve this limitation by instead using \bemphdiscrete embeddings and controlling the number of bits in each dimension via hyperparameters. To this end, we propose a model that produces Discrete InfoMax COdes (DIMCO) via an end-to-end learnable neural network encoder.

This work’s primary contributions are as follows. We motivate mutual information as an objective for learning embeddings, and propose an efficient method of estimating it in the discrete case. We experimentally demonstrate that learned discrete embeddings are more memory- and time- efficient compared to continuous embeddings. Our experiments also show that using discrete embeddings helps meta-generalization by acting as an information bottleneck. We also provide theoretical support for this connection through an information-theoretic PAC bound that shows the generalization characteristics of learned discrete codes.

This paper is organized as follows. We propose our model for learning discrete codes in Section 2. We justify our loss function and also provide generalization bound for our setup in Section 3. In Section 5, we present experiments that We compare our method to related work in Section 4 and conclude our paper in Section 6.

2 Discrete Infomax Codes (DIMCO)

We present our model which produces Discrete InfoMax COdes (DIMCO). A deep neural network is trained end-to-end to learn $k$-way $d$-dimensional discrete codes that maximally preserves the information on labels. We outline the training procedure in Algorithm 1, and we also illustrate the overall structure in the case of $4$-way $3$-dimensional codes ($k=4,d=3$) in Figure 1.

2.1 Learnable Discrete Codes

Suppose that we are given a set of labeled examples, which are realizations of random variables $(X,Y)\sim p(\mathbf{x},y)$, where $X$ is the input, and its corresponding label is $Y$. We denote its codebook by $\widetilde{X}$, which is a compressed representation of $X$. Realizations of $X$ and $Y$ are denoted by $\mathbf{x}\in\mathbb{R}^{D}$ and $y\in\{1,\ldots,c\}$.

We construct a probabilistic encoder $p(\mathbf{\widetilde{x}}|\mathbf{x})$, which is implemented by a deep neural network, that maps an input $\mathbf{x}$ to a \bemph$\bm{k}$-way $\bm{d}$-dimensional code $\mathbf{\widetilde{x}}\in\{1,2,\ldots,k\}^{d}$. That is, each entry of $\mathbf{\widetilde{x}}$ takes on one of $k$ possible values and the cardinality of $\widetilde{X}$ is $|\widetilde{X}|=k^{d}$. Special cases of this coding scheme include $k$-way class labels ($d=1$), $d$-dimensional binary codes ($k=2$), and even fixed-length decimal integers ($k=10$).

We now describe our model which produces discrete infomax codes. A neural network encoder $f(\mathbf{x};\theta)$ outputs $k$-dimensional categorical distributions, $\mathrm{Cat}(p_{i,1},\ldots,p_{i,k})$. Here, $p_{i,j}(\mathbf{x})$ represents the probability that output variable $i$ takes on value $j$, consuming $\mathbf{x}$ as an input, for $i=1,\ldots,d$ and $j=1,\ldots,k$. The encoder takes $\mathbf{x}$ as an input to produce logits $\{l_{i,j}\}$, which are reshaped into a $d\times k$ matrix:

[TABLE]

These logits undergo softmax functions to yield

[TABLE]

Each example $\mathbf{x}$ in the training set is assigned a codeword $\mathbf{\widetilde{x}}=[\widetilde{x}_{1},\ldots,\widetilde{x}_{d}]^{\top}$, each entry of which is determined by one of $k$ events that is most probable, i.e.,

[TABLE]

While the stochastic encoder $p(\mathbf{\widetilde{x}}|\mathbf{x})$ indues a soft partitioning of input data, codewords assigned by the rule in (3) yields a hard partitioning of $X$.

2.2 Loss Function

The $i$-th symbol is assumed to be sampled from the resulting categorical distribution $\mathrm{Cat}(p_{i1},\ldots,p_{ik})$. We denote the resulting distribution over codes as $\widetilde{X}$, and a code as $\mathbf{\widetilde{x}}$. Instead of sampling $\mathbf{\widetilde{x}}\sim\widetilde{X}$ during training, we use a loss function that optimizes the expected performance of the entire distribution $\widetilde{X}$.

We train the encoder by maximizing the \bemphmutual information between the distributions of codes $\widetilde{X}$ and labels $Y$. The mutual information is a symmetric quantity that measures the amount of information shared between two random variables. It is defined as

[TABLE]

Since $\widetilde{X}$ and $Y$ are discrete, their mutual information is bounded from both above and below as $0\leq I(\widetilde{X};Y)\leq\log|\widetilde{X}|=d\log k$.

To optimize the mutual information, the encoder directly computes empirical estimates of the two terms on the right-hand side of (4). Note that both terms consist of entropies of categorical distributions, which has the general closed-form formula

[TABLE]

Let $\overline{p}_{ij}$ be the empirical average of $p_{ij}$ calculated using data points in a batch. Then, $\overline{p}_{ij}$ is an empirical estimate of the marginal distribution $p(\widetilde{x})$. We compute the empirical estimate of $H(\widetilde{X})$ by adding its entropy estimate for each dimension.

[TABLE]

We can also compute

[TABLE]

where $c$ is the number of classes. The marginal probability $p(Y=y)$ is the frequency of class $y$ in the minibatch, and $H(\widetilde{X}|Y=y)$ can be computed by computing (6) using only datapoints which belong to class $y$. We emphasize that such a closed-form estimation of $I(\widetilde{X};Y)$ is only possible because we are using discrete codes. If $\widetilde{X}$ was a continuous variable instead, we would have had to resort to approximtions of $I(\widetilde{X};Y)$ (e.g., Belghazi et al. (2018)) to optimize it.

We briefly examine the loss function (4) to see why maximizing it results in discriminative $\widetilde{X}$. Maximizing $H(\widetilde{X})$ encourages the distribution of all codes to be as dispersed as possible, while minimizing $H(\widetilde{X}|Y)$ encourages the average embedding of each class to be as concentrated as possible. Thus, the overall loss $I(\widetilde{X};Y)$ imposes a partitioning problem on the model: it learns to split the entire probability space into regions with minimal overlap between different classes. We will analyze our loss function more rigorously in Section 3.1.

2.3 Similarity Measure

Suppose that all data points in the training set are assigned their codewords according to the rule (3). Now we introduce how to compute a similarity between a query datapoint $\mathbf{x}^{(q)}$ and a support datapoint $\mathbf{x}^{(s)}$ so that it can be used for information retrieval or few-shot classification In the meta-learning literature, $\mathbf{x}^{(q)},\mathbf{x}^{(s)}$ are also called test data and query data, respectively (Finn et al., 2017).

Denote by $\mathbf{\widetilde{x}}^{(s)}$ the codeword associated with $\mathbf{x}^{(s)}$, constructed by (3). For the test data $\mathbf{x}^{(q)}$, the encoder yields $p_{i,j}(\mathbf{x}^{(q)})$ for $i=1,\ldots,d$ and $j=1,\ldots,k$. As a similarity measure between $\mathbf{x}^{(q)}$ and $\mathbf{x}^{(s)}$, we calculate the following log probability

[TABLE]

The probabilistic quantity (8) indicates that $\mathbf{x}^{(q)}$ and $\mathbf{x}^{(s)}$ become more similar when encoder’s output, when $\mathbf{x}^{(q)}$ is provided, is well aligned with $\mathbf{\widetilde{x}}^{(s)}$.

We can view our similarity measure (8) as a probabilistic generalization of the Hamming distance (Hamming, 1950). The Hamming distance quantifies the similarity between two strings of equal length as the number of positions at which the corresponding symbols are equal. Because we have access to a distribution over codes, we use (8) to directly compute the log probability of having the same symbol at each position.

We use (8) as a similarity metric for both few-shot classification and image retrieval. We perform few-shot classification by computing a codeword for each class via (3) and classifying each test image by choosing the class that has the highest value of (8). We similarly perform image retrieval by mapping each support image to its most likely code (3) and for each query image retrieving the support image that has the highest (8).

While we have described the operations in (3) and (8) for a single pair $(\mathbf{x}^{(q)},\mathbf{x}^{(s)})$, one can easily parallelize our evaluation procedure since it is an argmax followed by a sum111We show a parallel implementation in the supplementary material.. Furthermore, $\mathbf{\widetilde{x}}$ typically requires little memory as it consists of discrete values, allowing us to compare against large support sets in parallel. Experiments in Section 5.4 investigate the degree of DIMCO’s efficiency in terms of both time and memory.

2.4 Regularizing by Enforcing Independence

One way of interpreting the code distribution $\widetilde{X}$ is as a group of $d$ separate code distributions $\widetilde{x}_{1},\ldots,\widetilde{x}_{d}$. Note that the similarity measure described in (8) can be seen as ensemble of the similarity measures of these $d$ models. A classic result in ensemble learning is that using more diverse learners increases ensemble performance (Kuncheva & Whitaker, 2003). In a similar spirit, we used an optional regularizer which promotes pairwise independence between each pair in these $d$ codes. Using this regularizer stabilized training, especially in more large-scale problems.

Specifically, we randomly sample pairs of indices $i_{1},i_{2}$ from $\mathinner{\left\{1,\ldots,d\right\}}$ during each forward pass. Note that $\widetilde{x}_{i_{1}}\otimes\widetilde{x}_{i_{2}}$ and $(\widetilde{x}_{i_{1}},\widetilde{x}_{i_{2}})$ are both categorical distributions with support size $k^{2}$, and that we can estimate the two different distributions within each batch. We minimize their KL divergence to promote independence between these two distributions:

[TABLE]

We compute (9) for a fixed number of random pairs of indices for each batch. The cost of computing this regularization term is miniscule compared to that of other components such as feeding data through the encoder.

Using this regularizer in conjunction with the learning objective (4) yields the following regularized loss:

[TABLE]

We fix $\lambda=1$ in all experiments; we found that DIMCO’s performance was not particularly sensitive to this hyperparameter. We emphasize that while this optional regularizer stabilizes training, our learning objective is the mutual information $I(\widetilde{X};Y)$ in (4).

2.5 Visualization of Codes

In Figure 2, we show images retrieved using our similarity measure (8). We trained a DIMCO model ($k=16$, $d=4$) on the CIFAR100 dataset. We select specific code locations and plot the top $10$ test images according to our similarity measure. For example, the top-1 (leftmost) image for code $(\cdot,j_{2},\cdot,j_{4})$ would be computed as

[TABLE]

where $N$ is the number of test images.

We visualize two different combinations of codes in Figure 2. The two examples show that using codewords together results in their respective semantic concepts to be combined: (man + fish = man holding fish), (round + warm color = orange). While we visualized combinations of $2$ codewords for clarity, DIMCO itself uses a combination of $d$ such codewords. The regularizer described in Section 2.4 further encourages each of these $d$ codewords to represent different concepts. The combinatorially many ($k^{d}$) combinations in which DIMCO can assemble such codewords gives DIMCO sufficient expressive power to solve challenging tasks.

3 Analysis

3.1 Is Mutual Information a Good Objective?

Our learning objective for DIMCO (4) is the mutual information between codes and labels. In this subsection, we justify this choice by showing that many previous objectives are closely related to mutual information. Due to space constraints, we only show high-level connections here and provide a more detailed exposition in the supplementary material.

Cross-entropy

The de facto loss for classification is the cross-entropy loss, which is defined as

[TABLE]

where $\widehat{Y}$ is the model’s prediction of $Y$. Using the observation that the final layer $q(\cdot)$ acts a parameterized approximation to the true conditional distribution $p(Y|\widetilde{X})$, we write this as

[TABLE]

The $H(Y)$ term can be ignored since it is not affected by model parameters. Therefore, minimizing cross-entropy is approximately equivalent to maximizing mutual information. The two objectives become completely equivalent when the final linear layer $q(\cdot)$ perfectly represents the conditional distribution $q(\mathbf{y}|\mathbf{\widetilde{x}})$. Note that for discrete $\mathbf{\widetilde{x}}$, we cannot use a linear layer to parameterize $q(\mathbf{y}|\mathbf{\widetilde{x}})$, and therefore cannot directly optimize the cross-entropy loss. We can therefore view our loss as a necessary modification of the cross-entropy loss for our setup of using discrete embeddings.

Contrastive Losses

Many metric learning methods (Koch et al., 2015; Hoffer & Ailon, 2015; Sohn, 2016; Movshovitz-Attias et al., 2017; Duan et al., 2018) use a contrastive learning objective to learn a continuous embedding ($\widetilde{X}$). Such contrastive losses consist of (1) a positive term that encourages an embedding to move closer to that of other relevant embeddings and (2) a negative term that encourages it to move away from irrelevant embeddings. The positive term approximately minimizes $\log p(\widetilde{X}|y)$ while the negative term as approximately minimizes $\log p(\widetilde{X})$. Together, these terms have the combined effect of maximizing

[TABLE]

We show such equivalences in detail in the supplementary material.

In addition to these direct connections to previous loss functions, we show empirically in Section 5.1 that the mutual information strongly correlates with both the top-1 accuracy metric for classification and the Recall@1 metric for retrieval.

3.2 Does Using Discrete Codes Help Generalization?

In Section 1, we have motivated the use of discrete codes through the regularization effect of an information bottleneck. In this subsection, we prove a PAC bound to theoretically analyze whether learning discrete codes by maximizing mutual information leads to better generalization. In particular, we study how the mutual information on the test set is affected by the choice of input dataset structure and code hyperparameters $k,d$.

We analyze DIMCO’s characteristics by examining each minibatch. Following related meta-learning works (Amit & Meir, 2017; Ravi & Beatson, 2019), we call each batch a “task”. We emphasize that this is only a difference in naming convention. The analysis in this subsection applies equally well to the metric learning setup; we can view each batch consisting of support and query points as a task.

Define a task $T$ to be a distribution over ${\mathcal{Z}}={\mathcal{X}}\times{\mathcal{Y}}$. Let tasks $T^{1},\ldots,T^{n}$ be sampled i.i.d. from a distribution of tasks $\tau$. Each task $T$ consists of a fixed-size dataset $D_{T}=z_{T}^{1},\ldots,z_{T}^{m}=(x_{T}^{1},y_{T}^{1}),\ldots,(x_{T}^{m},y_{T}^{m})$, which is a set of $m$ i.i.d. samples from the data distribution ($z_{T}^{j}\sim T$). Let $\theta$ be the parameters of DIMCO. Let $X,Y,\widetilde{X}$ be the random variables for data, labels, and codes, respectively. Recall that our objective is the expected mutual information between labels and codes:

[TABLE]

The loss that we actually optimize (Equations 6 and 7) is the empirical loss

[TABLE]

The following theorem bounds the difference between these the expected loss ${\mathcal{L}}$ and the empirical loss $\hat{{\mathcal{L}}}$.

Theorem 1** (simplified).**

Let $d_{\Theta}$ be the VC dimension of the encoder $\widetilde{X}(\cdot)$. The following inequality holds with high probability:

[TABLE]

Proof.

We use VC dimension bounds and a finite sample bound for mutual information (Shamir et al., 2010). We defer a detailed statement and proof to the supplementary material. ∎

First note that all three terms in our generalization gap (29) converge to zero as $n,m\rightarrow\infty$. This shows that training a model by maximizing empirical mutual information as in Equations 6 and 7 generalizes perfectly in the limit of infinite data.

Theorem 1 also shows how the generalization gap is affected differently by dataset size $m$ and number of datasets $n$. A large $n$ directly compensates for using a large backbone ($d_{\Theta}$), while a large $m$ compensates for using a large final representation ($\widetilde{X}$). Put differently, to effectively learn from small datasets ($m$), one should use a small representation $(\widetilde{X})$. The number of datasets $n$ is typically less of a problem because the number of different ways to sample datasets is combinatorially large (e.g., $n>10^{10}$ for miniImagenet $5$-way $1$-shot tasks). Recall that DIMCO has $|\widetilde{X}|=d\log k$, meaning that we can control the latter two terms using our hyperparameters $d,k$. We have motivated the use of discrete codes through the information bottleneck effect of small codes $\widetilde{X}$, and Theorem 1 confirms this intuition.

4 Related Work

Information Bottleneck

DIMCO and Theorem 1 are both close in spirit to the information bottleneck (IB) principle (Tishby et al., 2000; Tishby & Zaslavsky, 2015; Shwartz-Ziv & Tishby, 2017). IB finds a set of compact representatives $\widetilde{X}$ while maintaining sufficient information about $Y$, minimizing the following objective function

[TABLE]

subject to $\sum_{\mathbf{\widetilde{x}}}p(\mathbf{\widetilde{x}}|\mathbf{x})=1$. Equivalently, it can be stated that one maximizes $I(\widetilde{X};Y)$ while simultaneously minimizing $I(\widetilde{X};X)$. Similarly, our objective (15) is information maximization $I(\widetilde{X};Y)$, while our bound (29) suggests that the representation capacity $|\widetilde{X}|$ should be low for generalization. In the deterministic information bottleneck (Strouse & Schwab, 2017), $I(\widetilde{X};X)$ is replaced by $H(\widetilde{X})$. These three approaches to generalization are related via the chain of inequalities $I(\widetilde{X};X)\leq H(\widetilde{X})\leq\log|\widetilde{X}|$, which is tight in the limit of $\widetilde{X}$ being imcompressible. For any finite representation, i.e., $|\widetilde{X}|=N$, the limit $\beta\rightarrow\infty$ in (18) yields a hard partitioning of $X$ into $N$ disjoint sets. DIMCO uses the infomax principle to learn $N=k^{d}$ such representatives, which are arranged by $k$-way $d$-dimensional discrete codes for compact representation with sufficient information on $Y$.

Meta-Regularization

Previous meta-learning methods have restricted task-specific learning by learning only a subset of the network (Lee & Choi, 2018), learning on a low-dimensional latent space (Rusu et al., 2018), learning on a meta-learned prior distribution of parameters (Kim et al., 2018), and learning context vectors instead of model parameters (Zintgraf et al., 2018). Our analysis in Theorem 1 suggests that reducing the expressive power of the task-specific learner has a meta-regularizing effect, indirectly giving theoretic support for previous works that benefitted from reducing the expressive power of task-specific learners.

Discrete Representations

Discrete representations have been thoroughly studied in information theory (Shannon, 1948). Recent deep learning methods directly learn discrete representations by learning generative models with discrete latent variables (Rolfe, 2016; van den Oord et al., 2017; Razavi et al., 2019) or maximizing the mutual information between representation and data (Hu et al., 2017). DIMCO is related to but differs from these works as it assumes a supervised meta-learning setting and performs infomax using labels instead of data.

A standard approach to learning label-aware discrete codes is to first learn continuous embeddings and then quantize it using an objective that maximally preserves its information (Gray & Neuhoff, 1998; Jegou et al., 2011; Gong et al., 2012). DIMCO can be seen as an end-to-end alternative to quantization which directly learns discrete codes. Jeong & Song (2018) similarly learns a sparse binary code in an end-to-end fashion by solving a minimum cost flow problem with respect to labels. Their method differs from DIMCO, which learns a dense discrete code by optimizing $I(\widetilde{X};Y)$, which we estimate with a closed-form formula.

Metric Learning

The structure and loss function of DIMCO is closely related to that of metric learning methods (Hoffer & Ailon, 2015; Sohn, 2016; Vinyals et al., 2016; Duan et al., 2018). We show that the loss functions of these methods can be seen as approximation to the mutual information ($I(\widetilde{X};Y)$) in Section 2.2, and provide a more in-depth exposition in the supplementary material. While all of these previous methods require a support/query split within each batch, DIMCO simply optimizes an information-theoretic quantity of each batch, removing the need for such structured batch construction.

Information Theory and Representation Learning

Many works have applied information-theoretic principles to unsupervised representation learning: to derive an objective for GANs to learn disentangled features (Chen et al., 2016), to analyze the evidence lower bound (ELBO) (Alemi et al., 2017; Chen et al., 2018), and to directly learn representations (Alemi et al., 2016; Hjelm et al., 2018; Oord et al., 2018; Grover & Ermon, 2018; Choi et al., 2019). DIMCO is also an information-theoretic representation learning method, but we instead assume a supervised learning setup where the representation must reflect ground-truth labels. We also used previous results from information theory to prove a generalization bound for our representation learning method.

5 Experiments

In our experiments, we use datasets with varying degrees of complexity: CIFAR10/100 (Krizhevsky, 2009), miniImageNet (Vinyals et al., 2016), CUB200 (Wah et al., 2011), Cars196 (Krause et al., 2013), and ImageNet (ILSVRC-2012-CLS, Deng et al. (2009)). We use standard train/test splits for each dataset unless stated otherwise. We also use various network architectures: 4-layer convnet (Vinyals et al., 2016) and ResNet12/20/50 (He et al., 2016; Mishra et al., 2017). We followed previously reported experimental setups as closely as possible, and provide minor experiment details in the supplementary material.

5.1 Correlation of Metrics

We have shown in Section 3.1 that the mutual information $I(\widetilde{X};Y)$ is strongly connected to previous loss functions for classification and retrieval. In this subsection, we perform experiments to verify whether $I(\widetilde{X};Y)$ is a good metric that quantitatively shows the quality of the representation $\widetilde{X}$. We trained DIMCO on the miniImageNet dataset with $k=d=64$ for $20$ epochs. We plot the pairwise correlations between five different metrics: ($5,10,20$)-way $1$-shot accuracy, $\textrm{Recall@}1$, and $I(\widetilde{X};Y)$. The results in Figure 3 show that all five metrics are very strongly correlated. We observed similar trends when training with loss functions other than $I(\widetilde{X};Y)$ as well; we show these experiments in the appendix due to space constraints.

5.2 Label-Aware Bit Compression

We applied DIMCO to compressing feature vectors of trained classifier networks. We obtained penultimate embeddings of ResNet20 networks each trained on CIFAR10 and CIFAR100. The two networks had top-1 accuracies of $91.65$ and $66.61$, respectively. We trained on embeddings for the train set of each dataset, and measured top-1 accuracy of the test set using the training set as support. We compare DIMCO to product quantization (PQ, Jegou et al. (2011)), which similarly compresses a given embededing to a $k$-way $d$-dimensional code. We compare the two methods in Table 2 with the same range of $k,d$ hyperparameters. We performed the same experiment on the larger ImageNet dataset with a ResNet50 network which had a top-1 accuracy of $76.00$. We compare DIMCO to both adaptive scalar quantization (SQ) and PQ in Table 2. We show extended experiments for all three datasets in the supplementary material.

The results in Table 2 and Table 2 demonstrate that DIMCO consistently outperforms PQ, and is especially efficient when $d$ is low. Furthermore, the ImageNet experiment (Table 2) shows that DIMCO even outperforms SQ, which has a much lower compression rate compared to the embedding sizes we consider for DIMCO. These results are likely due to DIMCO performing label-aware compression where it compresses the embedding while taking the label into account, whereas PQ and SQ only compress the embeddings themselves.

5.3 Few-shot Classification

We evaluated DIMCO’s few-shot classification performance on the miniImageNet dataset. We compare against the following previous works: Ravi & Larochelle (2016); Vinyals et al. (2016); Snell et al. (2017); Sung et al. (2018); Bertinetto et al. (2018); Liu et al. (2018); Ye et al. (2018); Lee et al. (2019). All methods use the standard four-layer convnet with $64$ filters per layer222 some methods used more filters; we used $64$ for fairness. . We use the data augmentation scheme proposed by Lee et al. (2019) and use balanced batches of $100$ images consisting of $10$ different classes. We evaluate on both $5$-way $1$-shot and $5$-way $5$-shot learning, and report $95\%$ confidence intervals of $1000$ random episodes on the test split.

Results are shown in Table 3, and we provide an extended table with an alternative backbone in the supplementary material. Table 3 shows that DIMCO outperforms previous works on the $5$-way $1$-shot benchmark. DIMCO’s $5$-way $5$-shot performance is relatively low, likely because the similarity metric (Section 2.3) handles support datapoints individually instead of aggregating them, similarly to Matching Nets (Vinyals et al., 2016). Additionally, other methods are explicitly trained to optimize $5$-shot performance, whereas DIMCO’s training procedure is the same regardless of task structure.

5.4 Image Retrieval

We conducted image retrieval experiments using two standard benchmark datasets: CUB-200-2011 and Cars-196. As baselines, we consider three widely adopted metric learning methods: Binomial Deviance (Yi et al., 2014), Triplet loss (Hoffer & Ailon, 2015), and Proxy-NCA (Movshovitz-Attias et al., 2017). The backbone for all methods was a ResNet-50 network pretrained on the ImageNet dataset. We trained DIMCO on various combinations of $(p,d)$, and set the embedding dimension of the baseline methods to 128. We measured the time per query for each method on a Xeon E5-2650 CPU without any parallelization. We note that computing the retrieval time using a parallel implementation would skew the results even more in favor of DIMCO since DIMCO’s evaluation is simply one memory access followed by a sum.

Results presented in Table 4 show that DIMCO outperforms all three baseline, and that the compact code of DIMCO takes roughly an order of magnitude less memory, and requires less query time as well. This experiment also demonstrates that discrete representations can outperform modern methods that use continuous embeddings, even on this relatively large-scale task. Additionally, this experiment shows that DIMCO can train using large backbones without significantly overfitting.

6 Conclusion

We introduced DIMCO, a model that learns a discrete representation of data by directly optimizing the mutual information with the label. To evaluate our initial intuition that shorter representations generalize better between tasks, we provided generalization bounds that get tighter as the representation gets shorter. Our thorough experiments demonstrated that DIMCO is effective at both compressing a continuous embedding, and also at learning a discrete embedding from scratch in an end-to-end manner. The discrete embeddings of DIMCO outperformed recent continuous methods while also being more efficient in terms of both memory and time. We believe the tradeoff between discrete and continuous embeddings is an exciting area for future research.

DIMCO was motivated by concepts such as the minimum description length (MDL) principle and the information bottleneck: compact task representations should have less room to overfit. Interestingly, Yin et al. (2019) reports that doing the opposite—regularizing the task-general parameters—prevents meta-overfitting by discouraging the meta-learning model from memorizing the given set of tasks. In future work, we will investigate the common principle underlying these seemingly contradictory approaches for a fuller understanding of meta-generalization.

Appendix A Previous Loss functions Are Approximations to Mutual Information

Cross-entropy Loss

The cross-entropy loss has directly been used for few-shot classification (Vinyals et al., 2016; Snell et al., 2017).

Let $q(\mathbf{y}|\mathbf{\widetilde{x}};\phi)$ be a parameterized prediction of $\mathbf{y}$ given $\mathbf{\widetilde{x}}$, which tries to approximate the true conditional distribution $q(\mathbf{y}|\mathbf{\widetilde{x}})$. Typically in a classification network, $\phi$ is the parameters of a learned projection matrix and $q(\cdot)$ is the final linear layer. The expected cross-entropy loss can be written as

[TABLE]

Assuming that the approximate distribution $q(\cdot)$ is sufficiently close to $p(\mathbf{y}|\mathbf{\widetilde{x}})$, minimizing (19) can be seen as

[TABLE]

where the last equality uses the fact that $H(Y)$ is independent of model parameters. Therefore, cross-entropy minimization is approximate maximization of the mutual information between representation $\widetilde{X}$ and labels $Y$.

The approximation is that we parameterized $q(\mathbf{y}|\mathbf{\widetilde{x}};\phi)$ as a linear projection. This structure cannot generalize to new classes because the parameters $\phi$ are specific to the labels $\mathbf{y}$ seen during training. For a model to generalize to unseen classes, one must amortize the learning of this approximate conditional distribution. (Vinyals et al., 2016; Snell et al., 2017) sidestepped this issue by using the embeddings for each class as $\phi$.

Triplet Loss

The Triplet loss (Hoffer & Ailon, 2015) is defined as

[TABLE]

where $\mathbf{\widetilde{x}}_{q},\mathbf{\widetilde{x}}_{p},\mathbf{\widetilde{x}}_{n}\in\mathbb{R}^{d}$ are the embedding vectors of query, positive, and negative images. Let $\mathbf{y}_{q}$ denote the label of the query data. Recall that the pdf function of a unit Gaussian is $\log N(\mathbf{\widetilde{x}}|\mu,1)=-c_{1}-c_{2}\mathinner{\!\left\lVert\mathbf{\widetilde{x}}-\mu\right\rVert}_{2}^{2},$ where $c_{1},c_{2}$ are constants. Let $p_{p}(\mathbf{\widetilde{x}})=N(\mathbf{\widetilde{x}}_{p},1)$ and $p_{n}(\mathbf{\widetilde{x}})=N(\mathbf{\widetilde{x}}_{n},1)$ be unit Gaussian distributions centered at $\mathbf{\widetilde{x}}_{p},\mathbf{\widetilde{x}}_{n}$ respectively. We have

[TABLE]

Two approximations were made in the process. We first assumed that the embedding distribution of images not in $\mathbf{y}_{q}$ is equal to the distribution of all embeddings. This is reasonable when each class only represents a small fraction of the full data. We also approximated the embedding distributions $p(\mathbf{\widetilde{x}}|\mathbf{y}),p(\mathbf{\widetilde{x}})$ with unit Gaussian distributions centered at single samples from each.

N-pair Loss

Multiclass $N$-pair loss (Sohn, 2016) was proposed as an alternative to Triplet loss. This loss function requires one positive embedding $\mathbf{\widetilde{x}}^{+}$ and multiple negative embeddings $\mathbf{\widetilde{x}}_{1},$ $\ldots,$ $\mathbf{\widetilde{x}}_{N-1}$, and takes the form

[TABLE]

This can be seen as the cross-entropy loss applied to $\textrm{softmax}(\mathbf{\widetilde{x}}^{\top}\mathbf{\widetilde{x}}^{+},\mathbf{\widetilde{x}}^{\top}\mathbf{\widetilde{x}}_{1},\ldots,\mathbf{\widetilde{x}}^{\top}\mathbf{\widetilde{x}}_{N-1})$.

Following the same logic as the cross-entropy loss, this is also an approximation to $I(\widetilde{X};Y)$. This objective should have less variance than Triplet loss since it approximates $p(\mathbf{\widetilde{x}})$ using more examples.

Adversarial Metric Learning

Deep Adversarial Metric Learning (Duan et al., 2018) tackles the problem of most negative exmples being uninformative by directly generating meaningful negative embeddings. This model employs a generator which takes as input the embeddings of anchor, positive, and negative images. The generator then outputs a ”synthetic negative” embedding that is hard to distinguish from a positive embedding while being close to the negative embedding.

This can be seen as optimizing

[TABLE]

by estimating $p(\mathbf{\widetilde{x}})$ using a generative network rather than directly from samples. Rather than modelling the marginal distribution $p(\mathbf{\widetilde{x}})$, this method conditionally models $p(\mathbf{\widetilde{x}};\mathbf{\widetilde{x}}_{q},\mathbf{\widetilde{x}}_{p},\mathbf{\widetilde{x}}_{n})$ so that $\mathbf{\widetilde{x}}$ is hard to distinguish from $\mathbf{\widetilde{x}}_{p}$ while sufficiently close to both $\mathbf{\widetilde{x}}_{q}$ and $\mathbf{\widetilde{x}}_{n}$.

Appendix B Proof of Theorem 1

We restate and prove our main theorem.

Theorem 1**.**

Let $d_{\Theta}$ be the VC dimension of the encoder $\widetilde{X}(\cdot)$. Let $\hat{I}(\widetilde{X}(X_{T},\theta);Y_{T})$ be the empirical estimate of the mutual information using finite dataset $D_{T}$, and define empirical loss as

[TABLE]

The following inequality holds with high probability:

[TABLE]

Proof.

We use the following lemma from (Shamir et al., 2010), which we restate using our notation.

Lemma 1**.**

Let $\widetilde{X}$ be a random mapping of $X$. Let $D$ be a sample of size $m$ drawn from the joint probability distribution $p(X,Y)$. Denote the empirical mutual information observed from $D$ between $\widetilde{X}$ and $Y$ as $\hat{I}(\widetilde{X};Y)$. For any $\delta\in(0,1)$, the following holds with probability at least $1-\delta$:

[TABLE]

We simplify this and plug in our specific quantities of interest ($\widetilde{X}(X_{T},\theta)$, $Y_{T}$):

[TABLE]

We similarly bound the error caused by estimating ${\mathcal{L}}$ with a finite number of tasks sampled from $\tau$. Denote the finite sample estimate of ${\mathcal{L}}$ as

[TABLE]

Let the mapping $X\mapsto\widetilde{X}$ be parameterized by $\theta\in\Theta$ and let this model have VC dimension $d_{\Theta}$. Using $d_{\Theta}$, we can state that with high probability,

[TABLE]

where $d_{\Theta}$ is the VC dimension of hypothesis class $\Theta$.

Combining equations (33, 31), we have with high probability

[TABLE]

∎

Appendix C Experiments

Parameterizing the Code Layer

Recall that each discrete code is parameterized by a $\mathbb{R}^{k\times d}$ matrix. A problem with a naive implementation of DIMCO is that simply using a linear layer that maps $\mathbb{R}^{D}$ to $\mathbb{R}^{k\times d}$ takes $Dkd$ parameters in that single layer. This can be prohibitively expensive for large embeddings, e.g. $d=4096$. We therefore parameterize this code layer as the product of two matrices, which sequentially map $\mathbb{R}^{D}\rightarrow\mathbb{R}^{n}\rightarrow\mathbb{R}^{k\times d}$. The total number of parameters required for this is $nD+nkd$. We fix all $n=128$. While more complicated tricks could reduce the parameter count even further, we found that this simple structure was sufficient to produce the results in this paper.

Correlation of Metrics

We collect statistics from $8$ different independent runs, and report the average of $500$ batches of $1$-shot accuracies, Recall@1, and mutual information. $I(\widetilde{X};Y)$ was computed using balanced batches of $16$ images each from $5$ different classes. In addition to the experiment in the paper, we measured the correlation between $1$-shot accuracies, $\textrm{Recall}@1$, and NMI using three previously proposed losses (triplet, npair, protonet). Figure 4 shows that even for other methods for which mutual information is not the objective, mutual information strongly correlates with all other previous metrics.

Code Visualizations

We provide additional visualizations of codes in Figure 5. These examples consistently show that each code encodes a semantic concept, and that such concepts can be but are not necessarily tied to a particular class.

Label-aware Bit Compression

We computed top-1 accuracies using a kNN classifier on each type of embedding with $k=200$. We present extended results comparing PQ and DIMCO on ImageNet embeddings in Table 6. For CIFAR-10 and CIFAR-100 pretrained ResNet20, we used pretrained weights of open-sourced repository333https://github.com/chenyaofo/CIFAR-pretrained-models, and for ImageNet pretrained ResNet50, we used torchvision. We optimized the probablistic encoder with Adam optimizer (Kingma & Ba, 2014) with learning rate of 1e-2 for CIFAR-100 and ImageNet, and 3e-3 for CIFAR-10.

Small Train Set

We performed an experiment to see how well DIMCO can generalize to new datasets after training with a small number of datasets. We trained each model using $\{1,4,16,64\}$ samples from each class in the miniImageNet dataset. For example, $4$ samples means that we train on ($64$ classes\times$$4 images) instead of the full ($64$ classes\times$$600 images). We compare against three methods which use continuous embeddings for each datapoint: Triplet Nets (Hoffer & Ailon, 2015), multiclass N-pair loss (Sohn, 2016), and ProtoNets (Snell et al., 2017). We show the $\textrm{Recall}@1$ metric, and additionally show $\{5,10,20\}$-way $1$-shot accuracies in the supplementary material.

Figure 6 shows that DIMCO learns much more effectively compared to previous methods when the number of examples per class is low. We attribute this to DIMCO’s tight generalization gap. Since DIMCO uses fewer bits to describe each datapoint, the codes act as an implicit regularizer that helps generalization to unseen datasets. We additionally note that DIMCO is the only method in Figure 6 that can train using a dataset consisting of $1$ example per class. DIMCO has this capability because, unlike other methods, DIMCO requires no support/query (also called train/test) split and maximizes the mutual information within a given batch. In contrast, other methods require at least one support and one query example per class within each batch.

For this experiment, we used the Adam optimizer and performed a log-uniform hyperparameter sweep for learning rate $\in[\textrm{1e-7},\textrm{1e-3}]$ For DIMCO, we swept $p\in[32,128]$ and $d\in[16,32]$. For other methods, we made the embedding dimension $\in[16,32]$. For each combination of loss and number of training examples per class, we ran the experiment $64$ times and reported the mean and standard deviation of the top $5$.

Few-shot Classification

For this experiment, we built on the code released by Lee et al. (2019) (https://github.com/kjunelee/MetaOptNet) with minimal adjustments. We used the repository’s default datasets, augmentation, optimizer, and backbones. The only difference was our added module for outputting discrete codes. We show an extended table with citations in Table 7.