Set Prediction in the Latent Space

Set prediction tasks require the matching between predicted set and ground truth set in order to propagate the gradient signal . Recent works have performed this matching in the original feature space thus requiring predefined distance functions . We propose a method for learning the distance function by performing the matching in the latent space learned from encoding networks . This method enables the use of teacher forcing which was not possible previously since matching in the feature space must be computed after the entire output sequence is generated . Nonetheless , a naive implementation of latent set prediction might not converge due to permutation instability . To address this problem , we provide sufficient conditions for permutation stability which begets an algorithm to improve the overall model convergence . Experiments on several set prediction tasks , including image captioning and object detection , demonstrate the effectiveness of our method . Code is available at https : //github.com/phizaz/latent-set-prediction . 1 Introduction . Set prediction is a task where a model predicts multiple elements whose ordering is not relevant for correctness . This task is central to many real-world problems such as object detection , image captioning , and multi-speaker speech recognition . Object detection requires predicting a set of bounding boxes without any specific ordering . Describing objects within an image is a kind of image captioning yet perfectly suitable for set prediction . Multi-speaker speech recognition is also well suited for set prediction since the order of transcripts is irrelevant . Though these tasks can naturally be modeled as set prediction , traditional deep learning is not inherently suitable for these tasks . Multi-layer perceptrons and convolution networks with traditional loss functions impose a specific ordering on the prediction heads which hinders set prediction . A reasonable set prediction pipeline requires the model ’ s prediction heads to be more flexible . Each head does not have a predefined target , yet relies on its peers to determine what is best to predict to complete the target set . Recent works [ 1 , 2 ] emphasized using a Transformer model [ 3 ] , which is permutation-invariant , coupled with a permutation-invariant loss function as the main ingredients . Any traditional loss function can be made permutation-invariant by solving for a minimum bijective matching between predicted set and ground truth set via the Hungarian algorithm under a certain distance metric . After the assignment , the loss function is calculated between the assigned pairs , and backpropagation is performed accordingly . This scheme is known as Permutation Invariant Training ( PIT ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . A distance metric used by the assignment must agree with the loss function in a way that the assignment is kept after an optimization step on the loss function . A distance metric that fails this criterion may switch pairings hindering the convergence . Hence , a distance metric is crucial to the convergence property of the set optimization scheme . One may argue to use the loss function itself as a distance metric . However , not all loss functions have meaningful scalar values . For example , the vanilla GAN ’ s loss [ 4 ] is not insightful in terms of progress or distance . Combining loss functions from different domains also complicates the matter because they are not easily comparable in their scalar forms . In object detection , both L1 error loss and cross entropy loss are used to learn bounding box prediction [ 2 ] , but it is unclear how to define a proper distance metric from such a combination . Either hand-tuned coefficients or different surrogate distance metrics may be needed to form an effective distance metric . This begets the problem of selecting a proper distance metric for PIT . A set prediction scheme that does not require a hand-tuned distance function is appreciable . Another hardship related to PIT is when applying set prediction on sequence domains that require teacher forcing to train . Auto-regressive with teacher forcing is often used for sequence prediction such as speech recognition [ 5 , 6 ] and machine translation [ 7 , 3 ] . However , teacher forcing requires a groundtruth assignment before it can begin prediction . PIT also relies on the teacher forcing prediction to do minimum assignment , resulting in a chicken and egg problem . If the set cardinality is small enough , it is possible to exhaustively teacher force with respect to all possible ground truths requiring O ( N2 ) forward passes through the model , and keep only those with the minimum assignment distances for optimization . What if the Hungarian assignment is done in a latent space instead ? Since the latent space is learned , the choice of any specific distance metric is alleviated – even a simple Euclidean distance is reasonable . Since the latent space is prior to the sequence prediction , the prediction process knows exactly what its ground truth is which allows for efficient , O ( N ) , teacher forcing . This paper presents latent set prediction ( LSP ) which enables the assignment in the latent space with Euclidean distance metric . At the same time , it provides a convergence guarantee of the loss function by reducing the effect of permutation switches that can be problematic when performing matching in the latent space . Our contributions are as follows : 1 . We propose a framework for deep set prediction that alleviates the need for hand-crafted distance metrics . 2 . This framework is efficient for the set of sequence predictions with teacher forcing requiring only O ( N ) predictions , an improvement from the usual exhaustive O ( N2 ) . 3 . We provide a convergence proof of set prediction under this framework . 2 Related works . Set prediction . There are mainly two families of set prediction : distribution matching and minimum assignment . The distribution matching approaches learn P ( Y |x ) where Y is a set and x is an input . DeepSetNet and variants [ 8–10 ] proposed a likelihood function for set prediction . An energy function learned via adversarial samples was also proposed [ 11 ] . On the other hand , the minimum assignment approaches rely on solving assignment problems . The loss function is calculated between the assigned pairs afterward . Either Hungarian assignment ( bijection ) or Chamfer assignment is usually used depending on tasks . Zhang et al . [ 12 ] proposed to mold a primitive set into a target set via gradient signals from a set encoder . Kosiorek et al . [ 1 ] proposed a Transformer for set prediction . A similar kind of design was also used in end-to-end object detection [ 2 ] . Besides the two assignments , a stable marriage was proposed for set autoencoding pretraining [ 13 ] . Image captioning is not usually related to set prediction . This is true for impression captions such as MS-COCO [ 14 ] which only describe the most salient objects . A different kind of captioning is descriptive which describes individual objects in a scene and their interactions . A prominent example is Visual Genome [ 15 ] . For the same reason , a chest radiology report is also descriptive [ 16 ] . Since descriptive captions have no specific ordering , this task is actually a set prediction where the elements are captions themselves . To the best of our knowledge , there is no practical approach for set of text predictions that involves teacher forcing . Object detection is formulated as a set prediction task where each bounding box is a set member . Most object detection algorithms impose ordering by dividing the image into several grids . Each cell is responsible for predicting bounding boxes whose centers reside in it . Anchors are introduced to allow a single cell to support multiple objects [ 17 , 18 ] . As a result , hand-crafted components are needed for these methods to function efficiently . Though many works are focusing on removing the use of anchors [ 19 , 20 ] , the dense grid prediction still remains . Later , DETR [ 2 ] directly applies set prediction on bounding boxes whose process matches predicted boxes to ground truth boxes . To achieve a satisfactory result , the matching cost has to be manually designed . 3 Latent Set Prediction ( LSP ) . A common set prediction pipeline has three components : a set prediction model , ground truths , and an assignment mechanism . Our method is focused on the case of Hungarian assignment . Traditionally , the assignment mechanism matches the ground truths with the model predictions in an output space Y . Here , the pairing happens in a latent space RC . 3.1 Notations We assume a set prediction modelF : X → RN×C whereX is the input space andN is the cardinality of the set . Effectively , F outputs N vectors in a latent space RC . The model F is also responsible for set cardinality prediction . Each latent vector is passed through a decoder D : RC → Y where Y is any output space . Note that the decoder may also accept the input x ∈ X wherever the input is required for better prediction . To pair in the latent space , we utilizes another component called encoder E : Y ×X → RC where y ∈ Y is an output and x ∈ X is its corresponding input . The encoder maps elements of the output space as guiding vectors g ∈ RC in the latent space to facilitate the assignment . Given a set of latent vectors { s1 , s2 , . . . , sN } and guiding vectors { g1 , g2 , . . . , gN } , the minimum assignment π is π = argmin π′∈P N∑ i ‖sπ′ ( i ) − gi‖2 ( 1 ) where P is the set of all permutations of N letters . A switch is said to occur when sπ ( i ) changes after a gradient update as illustrated in Figure 2 . 3.2 Method . Latent set prediction ( LSP ) begins by feeding an input x into the set prediction model F resulting in a set of latent vectors s ’ s . s ’ s do not have designated targets until the corresponding guiding vector g ’ s are retrieved . To get g , each ground truth yi is mapped via the encoder resulting in gi . The assignment algorithm is performed between s ’ s and g ’ s resulting in the minimum bijective matching π . This pairs up each latent vector sπ ( i ) to the guiding vector gi and the associated ground truth yi . Knowing its target , s goes through the decoder D resulting in ŷ . The task loss function Ltask is calculated accordingly between ŷπ ( i ) and yi , and then the optimization is performed . Note that s receives no training signal from g ; g only gives s its goal . After an optimization step , as sπ ( i ) moves along the task gradient toward better prediction of yi , it may move away from gi and approach another guiding vector gj . This can cause a switch as demonstrated in Figure 2 . Our method incorporates several techniques that encourage stable pairing of s ’ s and g ’ s over time , which turns out to be sufficient for the convergence of Ltask . A pictorial description of the LSP framework is depicted in Figure 1 . The proof of convergence of Ltask ( Section 4 ) indicates that the convergence hinges on the ever smaller gaps between s ’ s and g ’ s under π . In fact , Ltask is bounded from above by a function of∑N i ∥∥sπ ( i ) − gi∥∥2 . Hence , not only that g ’ s give s ’ s their goals , g ’ s must also follow wherever s ’ s go . By closing the gaps , it is less likely for a switch to happen . Should a switch happen , it would only be between a short distance which is not as harmful to Ltask . This does not imply that we need to avert switches at all costs . We can simply assign gi to si for all i to guarantee no switches . However , it is ordered prediction , not set prediction . In a sense , switches should be welcomed as a sign of learning a natural ordering as long as in the long run the gaps are still closing . Therefore , we propose two mechanisms to make sure that the gaps between s ’ s and g ’ s are smaller over time . First , we enforce an asymmetric latent loss to bring s ’ s and g ’ s together : Ls→glatent = ∑ i 1 2 ‖sπ ( i ) − [ gi ] ‖22 Llatent = βLs→glatent + γL g→s latent Lg→slatent = ∑ i 1 2 ‖ [ sπ ( i ) ] − gi‖22 Ltotal = Llatent + Ltask ( 2 ) where [ · ] is stop gradient , and β , γ control the loss strengths . Ideally , we want to set β = 0 since s ’ s should only follow the training signals from Ltask . However , we found β = 0.1 to be useful in practice providing a bit of help for g ’ s to meet sπ ’ s . We set γ = 1 as the default value and found it to work well across experiments . However , the latent loss alone is not enough to guarantee convergence . This is because the latent loss can not anticipate the movement of s ’ s . Even for a pair of infinitesimally close s and g , any sizeable ∇sπLtask can break apart the two . The second part which completes the convergence proof is gradient cloning ( GC ) which copies the task gradient ∇sπLtask from sπ ’ s to their respective g ’ s . Theoretically , the distance between s and g is strictly decreasing which satisfies the requirement for convergence . In practice , the models that predict s ’ s and g ’ s may not be equally capable as one may go faster than the other . To allow for this discrepancy , we propose a stronger version of GC namely gradient cloning with rejection ( GCR ) . With GCR , the leader of each pair of sπ and g is slowed down when ‖∇sπLlatent −∇gLlatent‖2 is larger than d ‖∇Ltask‖2 . The constant d ( default d = 10−3 ) is indicating whether s and g are sufficiently far apart ( relative to ‖∇Ltask‖2 ) requiring a slower leader . The leader is slowed down by rejecting its ∇Ltask along the span of ∇Llatent . The one with an obtuse angle between its task and latent gradients is considered a leader : 〈∇sπLtask , ∇sπLlatent〉 < 0 ( in case of s ) or 〈∇sπLtask , ∇gLlatent〉 < 0 ( in case of g ) . d serves as a parame- ter for choosing between GC ( d =∞ ) and GCR with always rejection ( d = 0 ) . A smaller d puts a stronger tendency for converging s and g at the cost of slower learning of s. The idea is depicted in Figure 3 and described in Algorithm 2 . We summarize LSP in Algorithm 1 which can be implemented efficiently with modern deep learning frameworks . Algorithm 1 Single training step of Latent Set Prediction ( LSP ) Given x ∈ X , Y ∈ Y N , d ∈ R S← F ( x ) . S ∈ RN×C , latent set element prediction ( Inference only ) Ŷ ← D ( S ) . Ŷ ∈ Y N , prediction on output space ( Training only ) G← E ( Y , repeat ( x ) ) .G ∈ RN×C , ground truth encoding π ← Hungarian ( G , S ) . Equation 1 Ŷπ ← D ( Sπ ) . Ŷπ ∈ Y N , prediction on output space Llatent ← Llatent ( Sπ , G ) . Equation 2 Ltask ← Ltask ( Ŷπ , Y ) ∇SπLtask , ∇SπLlatent , ∇GLlatent ← Backprop ( Ltask + Llatent ) ∇Sπ ← GCR ( ∇SπLtask , ∇SπLlatent −∇GLlatent , d ) . Algorithm 2 ∇G ← GCR ( ∇SπLtask , ∇GLlatent −∇SπLlatent , d ) . Algorithm 2 Continue backpropagation to F , D , E ’ s parameters Algorithm 2 Gradient Cloning with Rejection ( GCR ) Given∇Ltask , ∇Llatent , and d ∈ R obtuse← 〈∇Ltask , ∇Llatent〉 < 0 . obtuse ∈ [ 0 , 1 ] N , obtuse angles indicate leading positions far← ‖∇Llatent‖2 > d · ‖∇Ltask‖2 . far ∈ [ 0 , 1 ] N , large latent gradients indicate large distances ∇̂Llatent ← Llatent‖Llatent‖2 ∇Ltask ← ∇Ltask − ( obtuse · far ) · ∇̂Llatent · 〈 ∇Ltask , ∇̂Llatent 〉 . Gradient rejection Return∇Ltask +∇Llatent

Summary: Previous set prediction networks have typically been learned by using a matching step (e.g. Hungarian algorithm, chamfer distance) post-output in order to match the predicted set with the target set before computing the loss and propagating the gradient. Because the sets are represented as vectors (one per set element) stacked up in matrices, the predicted and target sets have an arbitrary order that might not be compatible with each other to directly compare corresponding elements (that is there might be a permutation that must be applied to one of the sets in order to align the elements with the other).   This paper proposes to move the matching step inside the network during training, such that a learned representation of each target element (in the order provided by the target) is produced that can be matched with the latent representation of each predicted set element. As such, when the latent vectors corresponding to the input are decoded they are in an order that is compatible with the ordering of the target set.   The inherent problem that arises is that there are discontinuities when assignments between the target and predicted sets change, so this work proposes a neat workaround that allows the target encodings to move towards and 'chase' their matched predicted latent.