Are Disentangled Representations Helpful for Abstract Visual Reasoning?

In plain words

A good way to describe something is to give each of its separate qualities its own part of the description: one part for shape, one for size, one for colour, one for position. Researchers call this a disentangled representation. The idea is often said to help with later tasks, but there has been little careful evidence for it on hard tasks.

This paper tests the idea on puzzles like Raven's Progressive Matrices: a grid of pictures with one missing, where you must work out the rule that links the rows and pick the picture that fits. The authors build two such puzzle sets from two toy image datasets whose hidden qualities are known. They then train a large number of encoders that try to separate those qualities, and a reasoning model on top of each one. The main finding is that encoders which separate the qualities better lead to reasoning models that learn faster from fewer examples. The benefit is clearest early in training and fades once the models see many more examples.

TL;DR

• Asks whether unsupervised disentangled representations make abstract visual reasoning easier to learn §1.
• Builds two RPM-like puzzle sets from dSprites and 3dshapes, using the datasets' known factors of variation §3.
• Trains 360 disentanglement encoders across six model rows and 3600 Wild Relation Networks (WReN) on their embeddings §4.1 §4.2.2.
• Author-reported: FactorVAE and beta-VAE scores correlate with reasoning accuracy in the few-sample regime; reconstruction error correlates more after 100K steps §4.2.2.
• Use it as evidence for sample efficiency in a synthetic setting, not as a general claim; the code is Apache-2.0 and reusable.

Problem

Disentanglement is widely argued to help downstream tasks, but a systematic test on a complex downstream task was lacking; prior evidence came from simpler tasks such as single-factor classification §1 §2.

The disentanglement metrics used here need ground-truth factors of variation, so they cannot be computed on real unlabelled data. The paper therefore assumes a disentangled representation can be obtained and asks only whether it is useful once obtained §2.

Abstract visual reasoning is hard for standard networks. A WReN trained from a randomly initialised CNN often stalls, predicting all answers with equal probability, and its runs differ widely §4.2.1. This is why the paper compares representations against pretrained and ground-truth-factor inputs.

Method

Puzzle construction §3. Each instance is a 3×3 grid of panels with the bottom-right panel missing, plus six candidate answers, one correct fig. 1. The relation is always the AND type: one or more factors take the same value across each row. The number of constant factors (1, 2 or 3) is drawn uniformly. Wrong answers are made by resampling the non-fixed factors, and one fixed factor, until the rules fail §3. The paper does not name the sampling mode used; see Pitfalls.

Data sets §3 app. A:

• dSprites: the orientation factor is dropped. Background colour (5 grey shades) and object colour (6 hues) are added as factors. For the puzzles only, scale has 3 values and x and y positions have 4 values each.
• 3dshapes: all factors are kept. For the puzzles only, scale and azimuth take 4 values each.

Encoders §2 §4.1. Four unsupervised families are trained on the individual panels, each as an autoencoder with a regularised objective eq. 1:

• beta-VAE, with a fixed beta, and a variant with controlled capacity increase;
• FactorVAE, which penalises total correlation through a discriminator;
• beta-TCVAE, which penalises total correlation with a Monte Carlo estimate;
• DIP-VAE-I and DIP-VAE-II, which penalise moment mismatch of the aggregate posterior.

Each method uses six values of its main hyperparameter, five seeds and both datasets. That gives 6 × 6 × 5 × 2 = 360 encoders §4.1 app. A.

\mathbb{E}_{p(x)}[\mathbb{E}_{q_{\phi}(z|x)}[-\log p_{\theta}(x|z)]] + \lambda_1\mathbb{E}_{p(x)}[R_1(q_{\phi}(z|x))] + \lambda_{2}R_2(q_{\phi}(z)) eq. 1

The encoder output that parametrises q(z|x) is the representation §2. The paper does not say whether the mean or a sample is used; the code's example config points at a postprocessed/mean output (see Paper vs code).

Reasoning model §3 §4.2. The WReN scores each candidate answer against the eight context panels. Each answer gets its own pass through the same network, and the highest score wins:

\mathrm{WReN}(a, C) = f_{\phi}\Big(\sum_{e_1, e_2 \in E} g_{\theta}(e_1, e_2)\Big), \quad E = \{\mathrm{CNN}(c_1), \dots, \mathrm{CNN}(c_8)\} \cup \{\mathrm{CNN}(a)\} eq. 2

Here g is the edge MLP, applied to every pair of panel embeddings, and f is the graph MLP that turns the summed relations into one score §3. The reasoning model gets one of three inputs: a CNN learned from scratch, a frozen pretrained representation, or the ground-truth factor values §4.2.1.

# Score one candidate answer with a Wild Relation Network  [eq. 2] [§3]
# context: 8 panels, (8, H, W, C); answer: 1 panel, (H, W, C)
# Output: one scalar score for this answer

def wren_score(context, answer):
    # Embed all 9 panels with the same encoder  [§3]
    panels = concatenate(context, answer)     # (9, H, W, C)
    nodes = embed(panels)                     # (9, d)
    # Code-only step, not in the paper: append each node's one-hot
    # index before the edge MLP
    # [code: disentanglement_lib/evaluation/abstract_reasoning/models.py#L274]
    nodes = append_position_one_hot(nodes)    # (9, d + 9)
    # Apply the edge MLP g to every ordered pair of nodes, self-pairs
    # included  [eq. 2]
    # [code: evaluation/abstract_reasoning/relational_layers.py#L63-L72]
    edge_sum = 0
    for node_a in nodes:
        for node_b in nodes:
            edge_sum = edge_sum + g(concatenate(node_a, node_b))
    # Turn the summed relations into one score with the graph MLP f  [eq. 2]
    return f(edge_sum)

# Train one reasoning model on freshly generated puzzles  [§4.2] [app. A]
# Batch 32; 100 000 steps; Adam; evaluate every 1000 steps
# [code: config/abstract_reasoning_study_v1/stage2/example.gin#L35-L38]

def train_reasoner(model, make_puzzles, optimizer):
    for step in range(100000):
        # Draw a new batch, so no puzzle is seen twice  [§4.2]
        contexts, answers, labels = make_puzzles(batch_size=32)
        # One score per candidate, six candidates per puzzle  [§3]
        scores = stack([model(contexts, answer) for answer in answers])
        # Softmax cross-entropy over the six scores
        # [code: evaluation/abstract_reasoning/models.py#L100-L102]
        loss = cross_entropy(scores, labels)
        optimizer.step(loss)
        if step % 1000 == 0:
            # Accuracy on 100 batches of fresh puzzles  [§4.2]
            accuracy = evaluate(model, n_batches=100)

Analysis §4.2.2. For each encoder, the authors compute five disentanglement scores (BetaVAE, FactorVAE, MIG, DCI, SAP) and reconstruction error. They take the Spearman rank correlation between each score and WReN accuracy at 1K, 2K, 5K, 10K, 20K, 50K and 100K training steps.

Hyperparameters

Encoders (disentanglement models), shared unless noted app. A shared.gin:

Reasoning models (WReN) app. A §4.2 example.gin:

Puzzle generation §3: relations per row drawn uniformly from {1, 2, 3} constant factors; answers: six per puzzle.

Results

Author-reported. Setup: synthetic puzzles built from dSprites and 3dshapes §3; encoders are the 360 models above; reasoning models are WReN with the search space above §4.2. The paper gives no numeric baseline table. Headline results are stated in text and in figures; the correlation coefficients appear only as plots.

The paper's central claim is that more disentangled representations are more sample-efficient in this setting §4.2.2 §5.

Paper vs code

Code: disentanglement_lib at commit 86a644d4ed35c771560dc3360756363d35477357 (Apache-2.0). Linked from the paper's footnote §4.1, and confirmed by its README. The abstract reasoning code lives in disentanglement_lib/evaluation/abstract_reasoning/.

• Encoder kernel sizes. Paper: all four encoder convolutions use 4×4 kernels tab. 1. Code: the first two use 4×4, and the third and fourth use 2×2 architectures.py.
• DIP-VAE lambda_d. Paper: lambda_d is 10 for both DIP-VAE-I and II app. A. Code: lambda_d = lambda_d_factor × lambda_od, with factor 10 for DIP-I and 1 for DIP-II vae.py sweep.py. So for DIP-II the code uses lambda_d = lambda_od, while the paper states 10.
• DIP-VAE-II sweep name. Paper: sweeps beta for DIP-VAE-II app. A. Code: sweeps dip_vae.lambda_od for both variants sweep.py.
• Positional encoding. Code: appends a one-hot node index to each node embedding before the edge MLP models.py relational_layers.py. Paper: equation 2 has no such term eq. 2.
• Baseline CNN. Paper: the from-scratch CNN has the same architecture as the disentanglement encoders §4.2.1. Code: four 4×4 stride-2 convolutions (32, 32, 64, 64 filters, "same" padding) and a flatten, with no 256-unit dense layer and no latent output. The num_latent argument is accepted but not used in those layers models.py.
• Dropout value. Paper: "dropout of 0.25, 0.5 or 0.75" app. A. Code: passes 1 − value to the dropout layer, so the value is a keep fraction, not a drop fraction models.py.
• Representation statistic. Paper: the encoder output is the representation §2. The code's example config comment points at a postprocessed/mean output example.gin. The paper does not say whether it uses the mean.
• Sampling mode. Paper: the alternative answers are made by resampling factors until the rules fail §3. Code: easy and hard modes exist pgm_data.py; the paper's description resembles hard, but it does not name the mode. The study config does not set pgm.pgm_type, which is a required binding pgm_data.py; the test configs use easy1 cnn_embedder.gin.
• Self-pairs in the relation sum. Paper: the sum runs over e1, e2 in E eq. 2; it does not say whether e1 = e2 counts. Code includes all N × N ordered pairs, self-pairs included relational_layers.py.

Checked and matching: batch size 64, 10 latent dimensions, 300K steps, encoder Adam settings, the beta-VAE, FactorVAE, beta-TCVAE and DIP sweep values, and 5 seeds app. A sweep.py. The list above covers the differences found.

Pitfalls

• Figures carry the correlation results. The text gives no correlation values, so every numeric comparison in the disentanglement analysis is "not stated" in prose. Read them from the figures in §4.2.2 and the appendix.
• Averages hide large spread. Runs from scratch differ widely and can get stuck at uniform predictions §4.2.1. Use the averages over many runs, not single runs.
• The 80% and 90% figures are averages over models in the initial study §4.2.1. Do not read them as the accuracy of one method.
• Synthetic, one relation type. The puzzles use only the AND relation, with factors drawn from two toy datasets §3. Generalisation to other reasoning tasks is not tested §5.
• Quantisation chooses the factor values. The code keeps fewer values per factor by floor-ratio mapping pgm_data.py pgm_data.py. The paper gives only the reduced counts §3, so the exact subset of values an implementer picks is set by this mapping.
• Ground truth at evaluation time. The scores need ground-truth factors, so the paper's claim concerns usefulness once a representation exists, not how to find one §2.
• Cost. The footnote states about 2.73 GPU-years on NVIDIA P100 to reproduce §1. A smaller study needs fewer encoders and seeds; the paper does not say how results change with fewer.
• Metrics disagree in the few-sample regime. MIG and SAP correlate weakly §4.2.2; the paper suggests they struggle to separate partially disentangled models §4.2.2.
• Randomised search. The WReN space is sampled, with 10 configurations per encoder §4.2.2. Without the same draws, results will differ.
• Learning rate unset in the example config. The example config lists the options in a comment but does not bind a rate example.gin.

Lineage

• Builds on: Wild Relation Networks, Barrett et al. 2018, arXiv:1807.04225 (id taken from the code's docstring, models.py); relation networks, Santoro et al. 2017 §3 (no id given in the paper). The VAE objectives come from beta-VAE, FactorVAE, beta-TCVAE and DIP-VAE §2. The experimental setup follows Locatello et al. 2018 app. A.
• Compared against: a CNN learned from scratch, frozen pretrained representations, and ground-truth factor inputs §4.2.1; reconstruction error as an unsupervised baseline §4.2.2.
• Superseded by: not stated in this card's sources.