PMR: Prototypical Modal Rebalance for Multimodal Learning

In plain words

A model that learns from two kinds of input, such as sound and video, often learns mostly from the easier one. The other input is then used less than it could be, so the combined model falls short of what its inputs could give.

The method adds a simple "typical example" test for each input. For every class it keeps an average feature vector, and it checks how well each input's features land near the right class average. An input that is doing worse gets an extra training signal that pulls its features tighter around those class averages. Early in training, the strong input is also slowed down a little, so it does not take over before the weak one catches up.

The method does not change the network's shape or how the inputs are combined, so it can be added to many existing multimodal models.

TL;DR

• Per-class prototypes (mean feature vectors) act as nearest-prototype classifiers, one per modality, and measure which modality is lagging §4.1.
• The lagging modality gets an extra prototype cross-entropy loss; the dominant one gets an entropy regulariser in the first epochs §4.1 §4.2.
• Works with any fusion, since it only needs each modality's encoder output §4.1. Author-reported gains on three datasets and several fusions §5.3.
• The released code implements the prototype loss but not the entropy regulariser, and it uses a hard switch where the paper uses a clipped ramp (see Paper vs code).

Problem

Joint training of a multimodal network optimises one loss for all modalities. The authors argue that the better modality dominates the gradient and suppresses the learning of the weaker one, so the weaker branch under-uses its own features §1 §3. Existing remedies modulate the learning pace from the fused output, which is itself dominated by the better modality §1 §2.1, or they add extra branches or distillation §2.1. Their scope is also limited to particular fusion methods or architectures §1 §2.1.

The authors measure, on CREMA-D, that the gradient direction from the fused output drifts away from each uni-modal guidance direction during training fig. 2 §3. They conclude that scaling the gradient magnitude alone is not enough §3.

Method

Setup. Two modalities m0 and m1 have encoders phi0 and phi1 producing representations z0 and z1. A linear head on their concatenation gives the logits, f(x) = W[z0; z1] + b eq. 1. The standard cross-entropy is L_CE eq. 2. For diagnosis, the paper also scores each modality with its own linear head and with the fused output eq. 4 eq. 5.

Step 1: prototypes. The prototype of class k for modality i is the centroid of that modality's representations of the class's examples eq. 6. They are computed on a subset of the training data (one tenth) and updated once per epoch with a momentum rule §4.1 eq. 12 alg. 1.

Step 2: prototype classifier. Each modality is classified by a softmax over negative Euclidean distances to the prototypes eq. 7. The paper uses the probability of the true class, summed over a batch, as that modality's score. The imbalance ratio rho_t is the batch-level ratio of modality 0's score to modality 1's score eq. 8.

Step 3: prototype cross-entropy (PCE). For each modality, the loss is cross-entropy over the negative distances to its own prototypes eq. 9. It does not depend on the other modality's features §4.1.

Step 4: acceleration loss. The total is L_CE plus alpha times beta times PCE of modality 0, plus alpha times gamma times PCE of modality 1 eq. 10. The coefficients come from rho_t eq. 11:

• if rho_t < 1 (modality 0 is weaker): beta = clip(0, 1/rho_t - 1, 1), gamma = 0;
• if rho_t >= 1 (modality 1 is weaker): beta = 0, gamma = clip(0, rho_t - 1, 1).

clip(a, b, c) limits b to the interval [a, c] §4.1.

Step 5: prototype entropy regulariser (PER). During the first E_r epochs, the loss subtracts mu times gamma times the entropy H0 of a prototype softmax, and mu times beta times the entropy H1 of another one eq. 13. The coefficients are applied to the opposite modality from eq. 10, which the paper says slows the dominant modality early on §4.2. Read the indices with care (see Pitfalls).

Procedure. Algorithm 1 refreshes the prototypes at the start of each epoch from the subset, then for each mini-batch computes rho_t, selects the loss (with PER while e < E_r, otherwise eq. 10 only), and takes a gradient step alg. 1.

# Prototypical modal rebalance  [alg. 1]  [§4]
# Input:  paired batches (x0, x1, y); subset D_s; alpha, mu, eps, E, E_r
# Output: encoders phi0, phi1 and linear head W, b

def train(dataset, subset, E, E_r, alpha, mu, eps):
    # Initial prototypes from the subset  [§4.1] [eq. 6]
    protos0, protos1 = class_centroids(subset)
    for epoch in range(E):
        # Momentum refresh of prototypes once per epoch  [§4.1] [eq. 12]
        new0, new1 = class_centroids(subset)
        protos0 = eps * protos0 + (1 - eps) * new0
        protos1 = eps * protos1 + (1 - eps) * new1
        for x0, x1, y in batches(dataset):
            # Encode each modality and fuse for the linear head  [eq. 1]
            z0 = phi0(x0)                          # (B, d)
            z1 = phi1(x1)                          # (B, d)
            logits = W @ concat(z0, z1) + b        # (B, K)
            # Prototype class probabilities per modality  [eq. 7]
            p0 = softmax(-euclid(z0, protos0))     # (B, K)
            p1 = softmax(-euclid(z1, protos1))     # (B, K)
            # Batch imbalance ratio from true-class probabilities  [eq. 8]
            rho = sum(p0[y]) / sum(p1[y])
            # Prototype cross-entropy per modality  [eq. 9]
            pce0 = cross_entropy(-euclid(z0, protos0), y)
            pce1 = cross_entropy(-euclid(z1, protos1), y)
            # Coefficients from rho: the weaker modality is accelerated  [eq. 11]
            beta, gamma = modulation(rho)
            # Acceleration loss  [eq. 10]
            loss = cross_entropy(logits, y)
            loss = loss + alpha * beta * pce0 + alpha * gamma * pce1
            # Early-epoch entropy regulariser  [§4.2] [eq. 13]
            if epoch < E_r:
                loss = loss - mu * gamma * entropy(p0)
                loss = loss - mu * beta * entropy(p1)
            step(loss)


def modulation(rho):
    # Piecewise clipped ramp, as printed in eq. 11  [§4.1]
    if rho < 1:
        return clip(0, 1 / rho - 1, 1), 0
    return 0, clip(0, rho - 1, 1)

Notes on the pseudocode. The entropy lines follow the printed form of eq. 13 and the paper's indices are ambiguous; the released code omits them entirely (see Paper vs code).

Hyperparameters

Results

Author-reported. Setup: ResNet18 encoders on CREMA-D and AVE, a small CNN on CG-MNIST, one RTX 3090 §5.2. Accuracy in percent on the test split. Baseline means the same network without PMR. The paper reports no standard deviations in these tables.

Against other modulation methods on CREMA-D with concatenation: Modality-Drop 55.8, Gradient-Blending 56.7, OGM-GE 57.7, PMR without PER 60.3, PMR 61.1 tab. 2. On AVE with film, PMR 66.4 versus OGM-GE 65.1 tab. 2. Gated fusion on AVE drops slightly, from 63.5 to 62.7 tab. 1.

Ablation on the prototype subset: 1 percent of data gives 54.1 on CREMA-D and 62.1 on AVE; 10 percent gives 61.1 and 67.1; 100 percent gives 61.5 and 67.3 tab. 5.

Paper vs code

Official code: the repository fanyunfeng-bit/Modal-Imbalance-PMR at code_commit, MIT licence. The repository title matches the paper, and its training script contains local paths under the name "yunfeng" bash.sh. The paper itself does not link the repository, and the owner profile states no affiliation, so attribution rests on that name match. Search: the GitHub API repository search for "PMR prototypical modal rebalance" returned this single repository.

• Imbalance ratio. Matches: a ratio of batch sums of true-class probabilities eq. 8 main_Prototype_grad.py.
• Modulation coefficients. The paper uses a clipped ramp eq. 11. The code uses a hard switch: beta or lam equals alpha on one side and 0 on the other, with no ramp main_Prototype_grad.py. The two agree only where |rho - 1| is large enough for the clip to saturate at 1.
• Entropy regulariser (PER). Not implemented. The code has no entropy term and no mu eq. 13 main_Prototype_grad.py.
• Prototype loss. Matches: cross-entropy over negative Euclidean distances, with the true class as target eq. 9 main_Prototype_grad.py.
• Momentum. Equivalent to eq. 12 with eps = momentum_coef main_Prototype_grad.py. The paper does not state eps. The CREMA-D example script sets it to 0.0, so the prototypes are fully refreshed each epoch bash.sh.
• Prototype data source. The paper uses a subset of training data §5.2. The code uses the validation split for AVE and CG-MNIST [code: main_Prototype_grad.py#L534-L537] main_Prototype_grad.py, and a shuffled tenth of training batches otherwise main_Prototype_grad.py.
• Learning-rate schedule. Paper: gradual decay from 1e-3 to 1e-4 §5.2. Code: step decay by 0.1 at epoch 70 main_Prototype_grad.py main_Prototype_grad.py.
• Alpha. Paper: 1 or 2 §5.2. Code default 1.0 main_Prototype_grad.py; the CREMA-D film example uses 3.0 bash.sh.
• Modulation window. Code applies the modulation for epochs 0 to 100 by default main_Prototype_grad.py main_Prototype_grad.py. The paper's algorithm applies it for all E epochs alg. 1.
• Visual feature. Code pools each visual sample over frames and space before the head, so the prototype features are the frame-averaged 512-d vector basic_model.py. The paper states no frame pooling for the prototypes.
• Unchecked. The MMTM and CentralNet runs of Table 3 and the AVE localisation runs of Table 4 were not traced into the code.

Pitfalls

• Name drift. The method is PMR in the abstract and most of the text, but Sections 5.3 and 5.4 say "PMA" §5.3 §5.4. Read these as PMR.
• Eq. 13 indices. The second entropy term is written with z0 and the modality-1 prototypes, which looks like a typo. The entropy argument is also a softmax over class distances, while the printed term selects the true class y. The paper does not clarify either point. Confirm the intended form before implementing PER; the code omits it.
• Epoch count and E_r. The total epochs and the number of PER epochs are not stated in the paper §5.2 alg. 1. The code default is 150 total [code: main_Prototype_grad.py#L37].
• Division by small sums. rho_t divides two batch sums and the code does the same without any epsilon main_Prototype_grad.py. A near-zero denominator makes 1/rho large before the clip. Add a guard before the clipping.
• Hard switch versus ramp. Implementing the printed clip with a binary switch, as the code does, changes the loss when rho is near 1.
• Subset size and quality. Table 5 shows that 1 percent of data can fall below the baseline on AVE tab. 5. The authors attribute this to a biased prototype §5.4.
• Momentum sensitivity. The paper does not give eps, and the example script's 0.0 differs from the code default 0.2. Results may depend on this.
• Cost. Prototype refresh needs a forward pass over the subset every epoch. The reference distance is a Python double loop main_Prototype_grad.py, so a vectorised version is needed for speed. The paper calls the prototype computation the main extra cost §5.4.
• Augmentation. Training crops visual frames to 224 with a random resized crop CramedDataset.py. The paper does not state this, and it changes the prototypes.
• Adaptive optimisers. The derivation assumes one learning rate for all parameters. The paper's AdaGrad and Adam results are empirical only §5.4 tab. 6.
• Statistics. Single numbers are reported without variance, and no seed count is stated tab. 1 tab. 2. Gaps of about 1 point, such as PMR with and without PER tab. 2, cannot be judged from the paper alone.

Lineage

• Builds on: prototypical networks, which use class centroids as non-parametric classifiers §2.2. Prior imbalance work named by the paper: OGM-GE, which modulates the dominant modality's gradient from the fused output §2.1 §3 (not carded; no content used).
• Compared against: Modality-Drop, Gradient-Blending and OGM-GE §5.3 tab. 2.
• Superseded by: not known from content in this wiki.