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: This paper proposes LSP, an approach to set-to-set prediction that doesn’t require handcrafting a distance metric. This distance metric is often needed to perform bijective matching (e.g., the Hungarian algorithm) between the predicted set and the ground-truth set. The key idea is to perform the matching in the latent space instead, which means that encoding each element of the ground-truth set to this latent space has to be performed. Further, the authors adds two “tricks” to help with convergence: asymmetric latent loss (Eq. (2) with hyperparameter beta and gamma)) and gradient cloning with and without rejection (GC & GCR, with hyperparamter d to control the degree of rejection.)  The paper also provides the convergence analysis and experimental validation on 1 synthetic and 3 real datasets.