Brain visualizations of teacher and student seem to differ quite a bit

[laplaceon]

Hello, I'm just wondering if this is expected behavior. I ran example inputs on both TribeV2 and TinyTribeV3 and plotted the brain visualizations.
I used the sintel demo from https://download.blender.org/durian/trailer/sintel_trailer-480p.mp4 and then this 10 second aladdin demo from https://drive.google.com/file/d/1rf2NBmyRMm8coWztd-wF40fFUtTwPkrJ/view?usp=drive_link.

Sintel

TribeV2

TinyTribeV3

Aladdin Trailer

TribeV2

TinyTribeV3

I checked out the Teacher vs. Student Comparison in the README and it seems the teacher and student can have some radical differences. I generated the student brain visualizations using both an ONNX export and then directly with the following code:

# ==============================================================================
# 0. DEPENDENCIES & ENVIRONMENT SETUP (Official Repo Pipeline)
# ==============================================================================
import os
import sys
import torch
import numpy as np
import warnings
from pathlib import Path
from google.colab import userdata, drive

# Authenticate Hugging Face
from huggingface_hub import login
hf_token = userdata.get('HF_TOKEN')
login(token=hf_token)
os.environ["HF_TOKEN"] = hf_token

# Mount Google Drive
drive.mount('/content/drive')
video_path = "/content/drive/MyDrive/TRIBE_Distillation_Project/sintel.mp4"

# Set up paths and import official rendering toolkit
warnings.filterwarnings('ignore')
sys.path.append('/content/tiny-tribe')

from tribev2.plotting import PlotBrain
from tiny_tribe.inference_kaggle import load_model_from_checkpoint, run_inference

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
TIMESTEPS = 5
SUBJECT_ID = 0

# ==============================================================================
# 1. INITIALIZE OFFICIALLY WRAPPED STUDENT PIPELINE
# ==============================================================================
print(f"⚙️ Initializing Verified PyTorch Pipeline on {DEVICE}...")

# Load the model using the repository's robust checkpoint wrapper
ckpt_path = "/content/tiny-tribe/checkpoints/best-epoch=052-val/pearson_r=0.7278.ckpt"
student_model = load_model_from_checkpoint(ckpt_path, DEVICE)

# ==============================================================================
# 2. EXTRACT FEATURES VIA OFFICIAL REPO PIPELINE
# ==============================================================================
print("📐 Extracting raw features using repository backbones...")
from tiny_tribe.inference_kaggle import build_input_from_video

# This single native call handles MobileViT, Whisper, text zero-padding, 
# and adds the required batch dimensions automatically!
student_inputs = build_input_from_video(video_path, DEVICE, n_frames=TIMESTEPS)

# ==============================================================================
# 3. EXECUTE STUDENT INFERENCE
# ==============================================================================
print("🚀 Pushing verified features through Tiny-TRIBE Fusion Core...")
student_preds_400 = run_inference(student_model, student_inputs, SUBJECT_ID, DEVICE)

# ==============================================================================
# 4. REVERSE PARCELLATION & RENDER VISUALIZATION
# ==============================================================================
def unpack_parcels(parcel_vals: np.ndarray, n_vertices: int = 20484) -> np.ndarray:
    """Map (400,) → (20484,) by reversing uniform-chunk parcellation."""
    n_parcels = parcel_vals.shape[0]
    chunk = n_vertices // n_parcels
    out = np.zeros(n_vertices, dtype=np.float32)
    for i in range(n_parcels):
        out[i * chunk:(i + 1) * chunk] = parcel_vals[i]
    out[n_parcels * chunk:] = parcel_vals[-1]
    return out

# Convert the (400, 5) output matrix into full 20,484 vertex space
student_preds_20k = np.stack([unpack_parcels(student_preds_400[:, t]) for t in range(TIMESTEPS)])

print("\n🎨 Rendering Student Model Map...")
plotter = PlotBrain(mesh="fsaverage5")

fig_student = plotter.plot_timesteps(
    student_preds_20k,
    cmap="fire",
    norm_percentile=99,
    vmin=0.6,
    alpha_cmap=(0, 0.2),
    show_stimuli=False
)
fig_student.suptitle("STUDENT MODEL (Tiny-TRIBE PyTorch) - Verified Repository Inference", fontsize=16, fontweight='bold', y=1.02)

Both the ONNX and PyTorch models output the same visualizations. I'm wondering if this is working as intended. Obviously distillation with such downscaling will lead to discrepancies but how can I tell if the student is working as expected.