Energy-Based Learning for Scene Graph Generation

In plain words

Scene graph generators look at a picture and list the objects in it and the relations between them, such as "man standing on rock". The usual training scores every object label and every relation label on its own, as if they had nothing to do with each other. Nothing then stops a model from predicting two relations that cannot both be true, or from falling back on vague words such as "on" because they are the most common.

The paper changes the training signal so that the whole graph is judged at once. A small network gives each pairing of an image and a candidate scene graph an energy score: a low score means the graph fits the image. Training pushes the score of the correct graph below the score of the best competing graph the model can find. That competitor is found by a noisy search over the graph's labels. Because the score sees the whole graph, an inconsistent combination of labels is penalised as a unit.

The idea is added as a loss to existing scene graph models, which keep their own architectures. The authors report better mean recall on Visual Genome and GQA, better recall on zero-shot and few-shot relation triplets, and larger gains for relations with few training examples.

TL;DR

Adds an energy-based loss on top of existing scene graph generators (VCTree, Motifs, IMP and VCTree-TDE on Visual Genome; Transformer, Motifs and IMP on GQA). An energy model scores (image graph, scene graph) pairs, SGLD sampling finds a low-energy competing graph, and the true graph's energy is pushed below it. Author-reported gains in mean Recall@K on both datasets, plus zero-shot and few-shot gains. The code does not match the paper's equations in several places (loss form, sampler, learning rate), and the text's percentages do not all reproduce from the paper's own table; see Paper vs code and Pitfalls.

Problem

Cross-entropy training of a scene graph model sums one loss per object and per relation, so the likelihood factorises into independent terms eq. 1 §1. Each relation loss ignores the other relations in the graph, so an impossible relation is penalised no more than a plausible but wrong one §1 fig. 1. Training is also dominated by frequent predicates, which biases predictions toward generic relations such as "on" §1 §4.2. The paper's setting assumes a pre-trained object detector supplies boxes and region features, and that a scene graph model supplies an initial prediction §3.1 §4.3. The structure term is presented as an inductive bias that lets the model learn from fewer labels §1.

Method

Prose overview, then numbered steps:

1. Detect and featurise. A Faster R-CNN with a ResNeXt-101-FPN backbone is pre-trained and its weights are frozen. Region features for the object boxes are taken with RoIAlign §3.1 §4.3.
2. Predict. A scene graph model M (VCTree, Motifs, IMP, VCTree-TDE or Transformer) predicts an initial scene graph SG^0 and is trained with its own task loss L_t §3.1 §3.3 §4.1.
3. Build two graphs. The image graph G_I has one node per object region, initialised with the region feature §3.3. The scene graph SG has object nodes (label rows of O) and directed relation edges (label rows of R) §3.3. The ground-truth pair (G_I+, SG+) is built from ground-truth boxes §3.3.
4. Score. The energy is E(G_I, SG) = MLP[ f(EGNN(SG)) ; g(GNN(G_I)) ], where f and g are gated pooling functions. The EGNN and the GNN each run three rounds of message passing §3.4 eq. 6 §4.3.
5. Find a low-energy competitor. Starting from SG^0, stochastic gradient Langevin dynamics (SGLD) updates the node and edge labels for τ = 20 iterations, with a step size of 1 and gradients clipped to [-0.01, 0.01]. Node and edge states are normalised to [0, 1] after each update §3.3 eq. 3 §4.3.
6. Loss. The energy loss compares the ground-truth energy with the lowest energy found by SGLD eq. 2. A squared-energy term keeps energies from growing without bound eq. 4. The total loss adds the base model's task loss with relative weights λ_e, λ_r and λ_t, whose values are not stated eq. 5 §3.3.

The paper's equations, numbered as in the LaTeX source §3.3:

\log p(SG \mid I) = \sum_{i \in O} \log p(o_i \mid I) + \sum_{j \in R} \log p(r_j \mid I) eq. 1

\mathcal{L}_e = E_\theta(G_I^{+}, SG^{+}) - \min_{SG \in \mathcal{SG}} E_\theta(G_I, SG) eq. 2

O^{\tau+1} = O^{\tau} - \frac{\lambda}{2} \nabla_{O} E_\theta(G_I, SG^{\tau}) + \epsilon^{\tau}

R^{\tau+1} = R^{\tau} - \frac{\lambda}{2} \nabla_{R} E_\theta(G_I, SG^{\tau}) + \epsilon^{\tau}, \quad \epsilon^{\tau} \sim \mathcal{N}(0, \lambda) eq. 3

\mathcal{L}_r = E_\theta(G_I^{+}, SG^{+})^2 + E_\theta(G_I, SG)^2 eq. 4

\mathcal{L}_{total} = \lambda_e \mathcal{L}_e + \lambda_r \mathcal{L}_r + \lambda_t \mathcal{L}_t eq. 5

E_\theta(G_I, SG) = \mathrm{MLP}\left[ f(\mathrm{EGNN}(SG)) \;;\; g(\mathrm{GNN}(G_I)) \right] eq. 6

\mathbf{m}_i^{t} = \alpha \, \mathbf{W}_{nn}\Big(\sum_{j \in \mathcal{N}_i} \mathbf{n}_j^{t-1}\Big) + (1-\alpha)\, \mathbf{W}_{en}\Big(\sum_{j \in \mathcal{N}_i} \mathbf{e}_{j \rightarrow i}^{t-1}\Big), \quad 0 \le \alpha \le 1 eq. 7

\mathbf{d}_{i \rightarrow j}^{t} = \mathbf{W}_{ee}\,[\mathbf{n}_i^{t-1} \,\Vert\, \mathbf{n}_j^{t-1}] eq. 8

\mathbf{N} = \sum_{k} f_{gate}(\mathbf{n}_k) \odot \mathbf{n}_k eq. 9

\mathbf{E} = \sum_{ij} g_{gate}(\mathbf{e}_{i \rightarrow j}) \odot \mathbf{e}_{i \rightarrow j} eq. 10

Training step as the paper formulates it. The paper gives no algorithm box, so this is a reconstruction from §3.3 and eq. 2 to eq. 5; the lam_* weights are not given in the paper.

# Energy-based training step, paper formulation  [§3.3] [eq. 2] [eq. 4]
# Input:  image I, ground-truth scene graph sg_plus (with boxes)
# Output: scalar loss

def energy_training_step(image, sg_plus):
    # Ground-truth image graph from its object boxes  [§3.3]
    image_graph_plus = region_graph(image, sg_plus.boxes)
    # Initial scene graph and task loss from the base model  [§3.1]
    sg_zero, task_loss = base_model(image)
    image_graph = region_graph(image, sg_zero.boxes)
    # Lowest-energy competitor, approximating the min in eq. 2  [eq. 3]
    sg_low = sgld(energy_model, image_graph, sg_zero, steps=20)
    # Energies of the ground truth and of the competitor  [§3.3]
    energy_plus = energy_model(image_graph_plus, sg_plus)
    energy_low = energy_model(image_graph, sg_low)
    # Energy loss: ground-truth energy minus the lowest energy  [eq. 2]
    loss_energy = energy_plus - energy_low
    # Squared-energy term against divergence  [eq. 4]
    loss_l2 = energy_plus**2 + energy_low**2
    # Weighted sum; lam_* are not stated in the paper  [eq. 5]
    return lam_e * loss_energy + lam_r * loss_l2 + lam_t * task_loss

SGLD over the scene graph. The code's step size and noise are separate settings, which differs from the paper's single λ (see Paper vs code).

# Sampler over the scene graph (SGLD)  [§3.3] [eq. 3] [§4.3]
# Input:  energy model, image graph, initial scene graph, number of steps
# Output: scene graph with (approximately) lower energy

def sgld(energy_model, image_graph, sg, steps):
    for step in range(steps):
        # Gradient of the energy w.r.t. node and edge states  [eq. 3]
        grad_nodes, grad_edges = energy_gradients(energy_model,
                                                  image_graph, sg)
        # Clip every gradient entry to [-0.01, 0.01]  [§4.3]
        grad_nodes = clip(grad_nodes, -0.01, 0.01)
        grad_edges = clip(grad_edges, -0.01, 0.01)
        # Gradient step (size 1) plus Gaussian noise  [§4.3] [eq. 3]
        sg.nodes = sg.nodes - 1.0 * grad_nodes + gaussian_noise()
        sg.edges = sg.edges - 1.0 * grad_edges + gaussian_noise()
        # Map every state back to the range [0, 1]  [§4.3]
        sg.nodes = min_max_normalise(sg.nodes)
        sg.edges = min_max_normalise(sg.edges)
    return sg

Energy model forward pass. Edge-aware messages are described in eq. 7 and eq. 8.

# Energy of an (image graph, scene graph) pair  [§3.4] [eq. 6]
# Input:  image graph (one node per region), scene graph (objects, edges)
# Output: scalar energy

def energy(image_graph, scene_graph):
    # Refine scene graph states with edge-aware messages  [§3.4.1]
    nodes, edges = edge_gnn(scene_graph, rounds=3)
    # Refine image graph states with a plain graph network  [§3.4]
    image_nodes = graph_net(image_graph, rounds=3)
    # Gated pooling of each graph into one vector  [§3.4.2] [eq. 9]
    scene_vector = gated_pool(nodes, edges)
    image_vector = gated_pool(image_nodes)
    # Concatenate both vectors and map them to one number  [eq. 6]
    return mlp(concat(scene_vector, image_vector))

Pointers for an implementer: the repository is developed on top of the Scene-Graph-Benchmark framework of Tang et al. README.md, and its LICENSE says it is modified from maskrcnn-benchmark LICENSE. The training entry point is tools/energy_joint_train_cd.py, which trains the task loss and the energy loss together README.md energy_joint_train_cd.py.

Hyperparameters

Values come from the paper unless marked "code". Code paths are relative to the repository root at code_commit.

Results

Author-reported. Setup: Visual Genome, 108k images with 150 object and 50 predicate classes, 70/30 train/test split with 5k training images held out for validation §4. GQA, 1704 object and 311 relation classes, same split scheme §4. Metric: mean Recall@K (mR@K) for PredCls, SGCls and SGDet §4.2. Each base model is trained with cross-entropy and with the energy loss under an identical setup §4.3. Hardware is not stated.

The paper's headline claim is improvement of up to 21% on Visual Genome and 27% on GQA over cross-entropy training, in mean Recall@K [abstract] §1. The text states that VCTree gains 12.7%, 22.3% and 5.6% on mR@20 for PredCls, SGCls and SGDet, and that Transformer gains 8.5% (PredCls) and 25.92% (SGCls) on GQA §5. These percentages do not all reproduce from tab. 1; see Pitfalls.

Mean Recall@K, cross-entropy (CE) vs energy-based loss (EBM-Loss) tab. 1

Caption of tab. 1: quantitative results, mean Recall@K under the three settings. SGDet is omitted on GQA because the underlying object detector has low mAP §5.

Zero-shot recall, zsR@20/50 (CE vs EB) tab. 2

Caption of tab. 2: zero-shot recall, models trained with CE and with the energy-based loss (EB), on VG and GQA. Zero-shot triplets are those never seen in training §4.2.

Few-shot Recall@20, VCTree tab. 3

Caption of tab. 3: few-shot Recall@20 for VCTree trained with CE and with the energy-based loss. The paper reports larger gains when data is scarce tab. 3.

Sentence-to-graph retrieval (S2GR), gallery of 1000 and 5000 images tab. 4

Caption of tab. 4: S2GR compares scene graphs retrieved from captions. The paper reports 5% to 23% relative gains for VCTree and Motif §5.

Ablation on VCTree, predicate classification tab. 5

Caption of tab. 5: effect of the number of SGLD steps τ and of removing the image graph. The 20-step row matches the VCTree EBM-Loss row of tab. 1. The paper reports that more steps give higher recall at extra compute, and that removing the image graph lowers recall §5.

Other claims, from text only: the paper reports that energy-based models gain more on relations with fewer training annotations than their CE counterparts. It cites supplementary material for per-relation recall, which is not in the source archive §5. Qualitative examples are in fig. 3, showing more specific predicates (for example "standing on" in place of "on") and zero-shot triplets §5.

Paper vs code

Official repository: https://github.com/mods333/energy-based-scene-graph, linked from the paper's abstract footnote §1. Pinned at code_commit d522c2b7a360b7c84439960fed22c497b73eb168. Code paths below are relative to the repository root.

• Loss form. The paper's energy loss is E(G+, SG+) minus the minimum over SG of E(G, SG) eq. 2. The default loss is ContrastiveDivergence defaults.py: a batch mean of the positive energies minus a batch mean of the sampled negative energies, plus the sum of squared energies losses.py. It uses one sampled negative and no minimum. The README command does not change the default. The code also has SoftPlus and LogSumExp losses losses.py that are not in the paper.
• Loss weights. The paper leaves λ_e, λ_r and λ_t unstated eq. 5. The joint script adds the task loss and the energy loss with weight 1 energy_joint_train_cd.py. The config key ENERGY_MODEL.L2COEFF is defined defaults.py but never read in the repository. The same is true of ENERGY_MODEL.TRAINIG_MODE defaults.py.
• SGLD step and noise. Eq. 3 uses one λ for the step (λ/2) and the noise variance (N(0, λ)) eq. 3. The code has two settings: SAMPLER.LR for the step and SAMPLER.VAR as the standard deviation of the noise sampler.py sampler.py. In each iteration the noise is added before the gradient step sampler.py. The paper's "step size of 1" §4.3 matches the README value of SAMPLER.LR README.md, but the code default is 0.1 defaults.py. The default iteration count is 10 defaults.py, against 20 in the paper §4.3 and the README README.md.
• Which states SGLD updates. Eq. 3 updates object and relation states in all cases eq. 3. In PredCls the code updates only the edge states sampler.py. In SGCls and SGDet it updates node and edge states sampler.py.
• Learning rate. The paper uses SGD with an initial learning rate of 1e-2 §4.3. The README sets SOLVER.BASE_LR to 0.001 README.md, and the code default is 0.002 defaults.py. The optimizer multiplies the base learning rate by IMS_PER_BATCH energy_joint_train_cd.py build.py. The README uses 16 README.md, so the effective rate is 0.016, not 0.01.
• Config values the paper does not state. The baseline config sets BASE_LR 0.01 e2e_relation_X_101_32_8_FPN_1x.yaml, weight decay 1e-4 e2e_relation_X_101_32_8_FPN_1x.yaml, 40000 iterations e2e_relation_X_101_32_8_FPN_1x.yaml and WarmupReduceLROnPlateau e2e_relation_X_101_32_8_FPN_1x.yaml. The README overrides BASE_LR and the iteration count README.md. The code default momentum is 0.9 defaults.py.
• Box information. The paper initialises object nodes from the label rows of O, and edges from R §3.3. The code concatenates a 128-d embedding of the box geometry (9 numbers) to each object label before the embedding layer model_ebm.py model_ebm.py. The paper does not mention this.
• Data noise. The code adds Gaussian noise with standard deviation DATA_NOISE_VAR = 1e-4 to the ground-truth and predicted states defaults.py utils.py utils.py utils.py. The paper does not mention this.
• Predicted and ground-truth graphs. The ground-truth scene graph contains only the annotated relation pairs, the non-zero entries of the relation matrix utils.py. The predicted graph contains the detector's candidate pairs utils.py. The paper does not say which pairs enter each graph.
• Normalisation of predictions. Predicted relation scores are min-max normalised per row before use utils.py utils.py. Predicted object states are normalised in SGCls and SGDet but not in PredCls utils.py. The paper does not mention either step.
• EGNN kernels. Eq. 7 uses single kernels W_nn and W_en. The code uses two-layer MLPs with ReLU layers.py. Eq. 8 uses a single W_ee. The code uses two 1x1 convolutions with ReLU over the concatenated endpoint states layers.py layers.py.
• Message normalisation. The code divides node-to-node messages by the node degree layers.py and sums edge-to-node messages without normalisation layers.py. Eq. 7 has no degree normalisation.
• α. Eq. 7 keeps α as a hyperparameter in [0, 1] and gives no value. The code fixes α = 0.5 layers.py.
• Gating and update. Both graph layers update their states with a GRU cell layers.py layers.py. The paper refers to a gating mechanism §3.4.1 without more detail.
• Pooling gates. Eq. 9 and eq. 10 use gate functions that map states to scalars. The code's gates are plain linear layers with no sigmoid or softmax pooling.py pooling.py. The modules htheta_node and htheta_edge are defined pooling.py but not used in the forward pass pooling.py pooling.py.
• Energy head. The energy is Linear, ReLU, Linear to one scalar model_ebm.py. The paper gives no depth or width for the MLP eq. 6.
• Frozen modules. Backbone, RPN and box head are frozen by setting requires_grad to false energy_joint_train_cd.py, which matches the paper's frozen detector §4.3. The relation head remains trainable.
• Validation. The call to the sampler inside run_energy_val appears commented out energy_joint_train_cd.py. It is not verified whether validation scores therefore skip SGLD.
• Second script. tools/energy_train_net.py is a second training entry point energy_train_net.py. The README uses the joint script README.md. This card does not compare the two.
• Released weights. The README links pre-trained VCTree checkpoints with and without the energy loss README.md. They were not downloaded or checked against Table 1.

Pitfalls

• The paper's loss is not the code's default loss. Eq. 2 takes the minimum energy over SG. The default code loss uses batch means over one sampled negative, with the squared-energy term folded into the same expression losses.py. Implement the paper's loss only with the paper's text in mind, and state which variant you use. The paper does not say which variant produced Table 1.
• Loss weights are missing. λ_e, λ_r and λ_t have no values eq. 5. The code uses 1 for all of them energy_joint_train_cd.py. The paper gives no evidence that other values were tried.
• One λ in the paper, two settings in the code. The paper ties step size and noise variance to one λ eq. 3. The code sets them independently. The noise setting named VAR is passed as a standard deviation sampler.py. Its value is 0.001 in the README and is not stated in the paper.
• Step size of 1 is ambiguous. The paper says "step size of 1" §4.3 but its update uses λ/2 eq. 3, and λ is never given.
• Cost of sampling. Each training step runs τ = 20 forward and backward passes of the energy model sampler.py §4.3. The paper says training time rises with the number of steps, but it gives no measurements §5.
• Division by zero in min-max normalisation. The normalisation subtracts the row minimum and then divides by the row maximum utils.py sampler.py. A row whose entries are all equal gives a zero divisor. The code has no guard. Add one if you reimplement this.
• Arithmetic in the text does not match the table. From tab. 1, VCTree PredCls mR@20 goes 13.07 to 14.2, which is +8.6%, not the 12.7% in the text. VCTree SGDet mR@20 goes 5.31 to 5.67, which is +6.8%, not 5.6%. Transformer GQA PredCls mR@20 goes 1.17 to 1.28, which is +9.4%, not 8.5%. The SGCls and the Transformer SGCls figures (22.3%, 25.92%) are close to the table. The paper does not say how the text figures were computed.
• The abstract's maxima do not match a stated cell. The abstract says "up to 21%" on Visual Genome and "27%" on GQA [abstract]. The paper does not name the cells. Table 1 contains larger relative gains than either figure: on GQA the Transformer's SGCls mR@100 goes 1.29 to 1.77, which is +37%. On Visual Genome IMP's SGDet mR@20 goes 2.20 to 2.78, which is +26%. Quote these numbers with the cell, not as a single headline.
• Ablation differences are small and unrepeated. Removing the image graph changes mR@20 from 14.20 to 14.19 tab. 5. The paper reports one run per setting and no variance. Do not treat the 0.01 gap as evidence.
• Few-shot table is underspecified. The text describes a "scene graph detection model" §5, but the caption says VCTree and names no setting tab. 3. The buckets 16-20 and 20-25 both contain 20 tab. 3.
• S2GR metric. The text says the paper reports Recall@20/100 §4.2, but tab. 4 reports R@20, R@50 and R@100.
• Frequency bias confounds the datasets. The bias is used on Visual Genome and dropped on GQA for memory reasons §4.3 §5. The two datasets are not directly comparable on this factor, and the GQA numbers are lower partly because of the larger label sets §5.
• Gradients from the energy loss are not stopped at region features. get_pred_im_graph calls detach() but discards the result, so the call does nothing utils.py. The relation head is trainable energy_joint_train_cd.py, so energy-loss gradients reach it. The paper does not say whether this is intended.
• Licence. The code's LICENSE is three lines: "MIT License" and an author note that says the license is undecided LICENSE. Per-file SPDX headers say MIT model_ebm.py. The repository says it is modified from maskrcnn-benchmark. Treat the code licence as unresolved. This card cites the code and does not quote it. The paper itself is CC BY 4.0.
• Dataset preparation. The Visual Genome split is a pre-processed version from [xu2017scene], and the GQA normalisation is described only in words §4. This card did not check the released data files against either description.

Lineage

• Builds on: the base scene graph models VCTree [tang2019learning], Neural Motifs [zellers2018neural], IMP [xu2017scene] and VCTree-TDE [tang2020unbiased]; the frequency-bias and mR@K evaluation from [tang2020unbiased]; SGLD [welling2011bayesian]; gated graph networks [DBLP:journals/corr/LiTBZ15]; energy-based learning [lecun2006tutorial]. The paper's bibliography gives no arXiv ids for these, and none has a card in this wiki yet.
• Compared against: the same base models trained with cross-entropy, throughout tab. 1 to tab. 4. The S2GR protocol follows [tang2020unbiased] §4.2.
• Superseded by: not known from content in this wiki.