# ImageDream: Image-Prompt Multi-view Diffusion for 3D Generation

Peng Wang Yichun Shi
ByteDance, USA
{peng.wang<EMAIL_ADDRESS>
###### Abstract

We introduce ImageDream, a novel Image-Prompt Multi-view diffusion model
devised for 3D object generation. ImageDream excels in delivering 3D models of
superior quality comparing with other State-of-the-Art (SoTA) image-
conditioned endeavors. Specifically, we consider a canonical camera
coordination of the object in image rather than relative. This enhancement
significantly augments the visual geometry correctness. Our models are
meticulously crafted, taking into account varying degrees of control
granularity derived from the provided image: wherein, the global control
predominantly influences the object layout, whereas the local control adeptly
refines the image appearance. The prowess of ImageDream is empirically
showcased through a comprehensive evaluation predicated on a common prompt
list as delineated in MVDream [37]. Our project page is https://Image-
Dream.github.io.

Figure 1:  ImageDream is a novel framework that generates high quality 3D
model from any viewpoint given a single image. It vastly improves the 3D
geometry quality comparing to previous SoTA, e.g. Magic123 [31], and more
importantly, it inherits the great text image alignment from the generated
image-prompt, comparing with MVDream [37]. Here, we provide 8 views of a
generated object from different methods, and in the last row, we show the
corresponding normal maps rendered with ImageDream generated model.

##  1 Introduction

Figure 2: The training pipeline of ImageDream. The blue arrow indicates
training of the diffusion network and the green arrow indicates training of
NeRF model. In diffusion training, given a 3D object, we first render multiple
views based on canonical camera coordination (bottom), and render another
image-prompt front-view images with a random setting (top). The multi-view
images are fed as training targets for multi-view diffusion networks, and
image-prompt is encoded with a multi-level controller as input to the
diffusion. In NeRF training, we use the trained diffusion for image-prompt
score distillation.

In the domain of 3D generation, incorporating images as an additional modality
for 3D generation, compared to methods relying solely on text [30], offers
significant advantages, as the common saying, An image is worth a thousand
words. Primarily, images convey rich, precise visual information that text
might ambiguously describe or entirely omit. For instance, subtle details like
textures, colors, and spatial relationships can be directly and unambiguously
captured in an image, whereas a text description might struggle to convey the
same level of detail comprehensively or might require excessively lengthy
descriptions. This visual specificity aids in generating more accurate and
detailed 3D models, as the system can directly reference actual visual cues
rather than interpret textual descriptions, which can vary greatly in detail
and subjectivity. Moreover, using images allows for a more intuitive and
direct way for users to communicate their desired outcomes, particularly for
those who may find it challenging to articulate their visions textually. This
multimodal approach, combining the richness of visual data with the contextual
depth of text, leads to a more robust, user-friendly, and efficient 3D
generation process, catering to a wider range of creative and practical
applications.

Adopting images as an additional modality for 3D object generation, while
beneficial, also introduces several challenges, Unlike text, images contain a
multitude of features like color, texture, spatial relationships that are more
complex to analyze and interpret accurately with a solely encoder like CLIP
[32]. In addition, high variant of light, shape or self-occlusion of the
object can lead to inaccurate and in-consistent view synthesis, therefore
leading blurry or incomplete 3D models.

The complexity of image processing necessitates advanced, computationally
intensive algorithms to accurately decode visual information and ensure
consistent appearance across multiple views. Researchers have employed various
strategies with diffusion models, such as Zero123 [19], and other recent works
[31, 20], to elevate a 2D object image to a 3D model. However, a limitation of
image-only solutions is that, although the synthesized views are visually
impressive, the reconstructed models often lack geometric accuracy and
detailed textures, particularly in the object’s rear views. This issue
primarily stems from significant geometric inconsistencies across the
generated or synthesized views. Consequently, during reconstruction, non-
matching pixels are averaged in the final 3D model, leading to indistinct
textures and smoothed geometry.

Fundamentally, image-conditioned 3D generation represents an optimization
problem with more stringent constraints compared to text-conditioned
generation. Hence, achieving optimized 3D models with clear details is more
challenging, as the optimization process is prone to deviating from the
trained distributions due to the limited amount of 3D data. For example,
generating a horse based solely on text descriptions may yield detailed models
if the training dataset includes a variety of horse styles. However, when an
image specifies particular textures, shapes, and fur details, the novel-view
texture generation may easily deviate from the trained distributions.

In this paper, we introduce ImageDream to address these challenges. Our
approach involves considering a canonical camera coordination across different
object instances and designing a multi-level image-prompt controller that can
be seamlessly integrated into the existing architecture. Specifically, the
canonical camera coordination mandates that the rendered image, under default
camera settings (i.e., identity rotation and zero translation), represents the
object’s centered front-view. This significantly simplifies the task of
mapping variations in the input image to 3D. The multi-level controller offers
hierarchical control, guiding the diffusion model from the image input to each
architectural block, thereby streamlining the path of information transfer.

As illustrated in Fig.1, ImageDream excels in generating objects with correct
geometry from a given image, enabling users to leverage well-developed image
generation models[29] for better image-text alignment than purely text-
conditioned models like MVDream [37]. Furthermore, ImageDream surpasses
existing state-of-the-art (SoTA) zero-shot single image 3D model generators,
such as Magic123 [31], in terms of geometry and texture quality. Our
comprehensive evaluation in the experimental section (Sec. 4), which includes
both qualitative comparisons through user studies and quantitative analyses,
demonstrates ImageDream’s superiority over other SoTA methods.

##  2 Related Works

Figure 3: The Multi-level Controller of ImageDream. Given an image prompt,
global controller and local controller takes input of image features after
CLIP encoding, and then output adapted features to cross-attention layers. It
represents image semantic information. The pixel controller send the VAE
encoded feature to diffusion, and perform pixel-level dense self-attention
with corresponding hidden features at each layer of the four-view MVDiffusion.

We recognize that 3D generation is a well-established field; this review
focuses on significant advancements closely related to our research.

Text-to-3D Generation with Diffusion. The emergence of deep generative models
has significantly impacted 3D generation. Early methods targeted the
reconstruction of simple objects using multi-view rendered images [11, 12].
The evolution of these techniques, from Generative Adversarial Networks
(GANs)[24, 25, 27, 9, 3, 10] to diffusion-based frameworks[33], marks a
notable progression.

Recent 3D diffusion models, specifically for tri-planes [38, 45] and feature
grids [15], have emerged. However, these models often focus on specific
objects like faces and ShapeNet [5] objects. Concurrently, there’s growing
interest in reconstructing object shapes from monocular image inputs [48, 26,
14], demonstrating the evolving stability of image generation methodologies. A
significant challenge remains in generalizing these models to the extent of
their 2D counterparts, likely due to constraints in 3D data size,
representation, and architectural design.

Lifting 2D Diffusion for 3D Generation. In light of the limited
generalizability of direct 3D generative models, a parallel line of research
has explored the elevation of 2D diffusion priors into 3D generation, often
integrating with 3D representations like NeRF [23]. A pivotal approach in this
area is the score distillation sampling (SDS) introduced by Poole et al.[30],
using diffusion priors as score functions to guide 3D representation
optimization. Alongside Dreamfusion, works like SJC[44], which utilize stable-
diffusion models [33], have emerged. Subsequent studies have focused on
enhancing 3D representations [17, 43, 42, 6], refining sampling schedules
[13], and optimizing loss designs [46]. Despite their ability to generate
photorealistic objects of various types without 3D data training, these
methods struggle with multi-view consistency. Moreover, each 3D model requires
individualized optimization through prompt and hyper-parameter adjustments.
Notably, MVDream [37] enhances generation robustness by joint training with 2D
and 3D datasets, producing satisfactory results with uniform parameters,
drawing on multi-view diffusion via SDS. Our work builds upon these concepts,
applying them to image-prompt generation and retaining the robustness
characteristic of MVDream.

Image-based Novel View Synthesis. Direct synthesis of novel 3D views from
single images has also been explored, bypassing traditional reconstruction
processes. Watson et al.[47] pioneered diffusion model applications in view
synthesis as the pipeline in Sitzmann et al.[39] using the ShapeNet dataset.
Subsequent advancements include Zhou et al.’s [53] extension to latent space
with an epipolar feature transformer and Chan et al.’s [4] approach to enhance
view consistency. Szymanowicz et al.[41] proposed a multi-view reconstructor
using unprojected feature grids. A common limitation across these methods is
their dependency on specific training data, with no established adaptability
to diverse image inputs. Fine-tuning pre-trained image diffusion models[2] on
extensive 3D render datasets for novel view synthesis, as proposed by Zero123
[19], remains constrained by geometric consistency issues. Later works,
including SyncDreamer [20], Consistent 1-to-3 [52], and Zero123plus [36], have
sought to enhance multi-view consistency through joint diffusion processes,
but the reconstruction of geometrically coherent 3D models remains a
challenge.

Single Image-conditioned Reconstruction. Recent advances in deriving 3D models
from single or few images predominantly leverage NeRF representations.
Techniques such as RegNeRF [28], which uses geometry loss from depth patches,
and SinNeRF [49], RealFusion [22], and NeuralLift [50], which combine depth
maps or Score Distillation Sampling during NeRF training, represent
significant steps forward. Despite their effectiveness, the quality of these
generated models remains suboptimal for real-world applications. Magic123 [31]
combines single-view and novel-view diffusion networks, achieving impressive
texture quality in 3D models. However, our tests reveal limitations in
understanding correct object geometry.

We also note recent parallel developments, such as Wonder3D [21], which
incorporate normal diffused outputs into original diffusion models, and
DreamCraft3D [40], which employ a second-stage DreamBooth-like model fine-
tuning for enhanced texture modeling. These works, while promising, remain
distinct from our contributions.

##  3 Methodology

In this section, we first talk about the MVDream [37] pipeline and then
describe our method to input the image prompt.

Figure 4: An example of the diffused results from different settings of multi-
level controllers in ImageDream (see Sec. 3.3).

###  3.1 Preliminary

In MVDream, there are two stage for 3D model production. The first stage is
training a multi-view diffusion network that produces four orthogonal and
consistent multi-view images from a text-prompt given respective camera
embedding. In the second stage, a multi-view score distillation sampling (MV-
SDS) is adopted to produce a detailed 3D NeRF model.

In the first stage, each block of the multi-view network contains a densely
connected 3D attention on the four view images, which allows a strong
interaction in learning the correspondence relationship between different
views. To train such a network, it adopts a joint training with the rendered
dataset from the Objaverse [8] and a larger scale text-to-image (t2i) dataset,
LAION5B [35], to maintain the generalizability of the fine-tuned model.
Formally, given text-image dataset $\mathcal{X}=\\{{\boldsymbol{x}},y\\}$ and
a multi-view dataset
$\mathcal{X}_{mv}=\\{{\boldsymbol{x}}_{mv},y,\boldsymbol{c}_{mv}\\}$, where
${\boldsymbol{x}}$ is an latent image embedding from VAE [16], $y$ is a text
embedding from CLIP [32], and $\boldsymbol{c}$ is their self-desgined camera
embedding, we may formulate the the multi-view (MV) diffusion loss as,

$\displaystyle\mathcal{L}_{\text{MV}}(\theta,\mathcal{X},\mathcal{X}_{mv})$
$\displaystyle=\mathbb{E}_{{\boldsymbol{x}},y,\mathbf{c},t,\epsilon}\left[\|\epsilon-\epsilon_{\theta}({\boldsymbol{x}}^{p};y,\boldsymbol{c}^{p},t)\|_{2}^{2}\right]$
(1) $\displaystyle\text{where,~{}}({\boldsymbol{x}}^{p},\boldsymbol{c}^{p})$
$\displaystyle=\begin{cases}({\boldsymbol{x}},\textbf{0})&\text{with
probability }p\\\ ({\boldsymbol{x}}_{mv},\boldsymbol{c}_{mv})&\text{with
probability }1-p\end{cases}$

here, ${\boldsymbol{x}}$ is the noisy latent image generated from a random
noise $\epsilon$ and image latent, the $\epsilon_{\theta}$ is the multi-view
diffusion (MVDiffusion) model parametrized by $\theta$.

After the model is trained, the MVDiffusion model can be inserted to the
DreamFusion [30] pipeline, where the authors adopt a score-distillation
sampling (SDS) based on the four generated views. Specifically, in each
iteration step, a random 4 orthogonal views are rendered from a NeRF $g(\phi)$
with a random 4 view camera extrinsic and intrinsic $\boldsymbol{c}$. Then,
they are encoded to latents ${\boldsymbol{x}}^{mv}$ and inserted to the multi
view diffusion network to compute a diffusion loss in the image space which is
back propagated to optimize the NeRF parameters. Formally,

$\mathcal{L}_{\text{MV-
SDS}}(\phi,{\boldsymbol{x}}^{mv})=\mathbb{E}_{t,\mathbf{c},\epsilon}\Big{[}\|{\boldsymbol{x}}^{mv}-\hat{{\boldsymbol{x}}}^{mv}_{0}\|_{2}^{2}\Big{]}.$
(2)

Here, $\hat{{\boldsymbol{x}}}^{mv}_{0}$ s the denoised MV image at timestep
$0$ from MVDiffusion. After fusion, MVDream shows significant improvement of
object geometry correctness without the Janus issues.

###  3.2 Canonical Camera

In the context of MVDream, a critical observation is the diffusion of multi-
view images using a global aligned camera coordination. In other words, the
image from a default camera (no azimuth rotation) is always the front view of
the object. This is done by asking the CLIP image feature of a view best match
the ”front view” CLIP text feature embeddings. This alignment facilitates the
fusion of diffused images in the fusion step, reducing ambiguity regarding
their viewpoints in relation to the provided text prompt.

As emphasized in the introduction, this alignment also reduced the
difficulties in learning the accurately reconstructing the geometry of
objects. Thereby, in image prompt cases, in contrast to previous image-
conditioned approaches like Zero123 [19], which attempt to recover object 3D
geometry based on image camera coordination system, ImageDream adopts
canonical/world camera coordination as in MVDream. Our diffusion model aims to
regress towards the canonical multiple view image of the object as depicted in
the image. This approach is expected to yield superior geometric accuracy
compared to systems that utilize relative camera coordination.

Formally, for an image of an object rendered from a random viewpoint with a
random camera, denoted as ${\boldsymbol{x}}_{r}$, we create the ImageDream
diffusion multi-view (MV) dataset as
$\mathcal{X}_{mv}=\\{{\boldsymbol{x}}_{mv},y,{\boldsymbol{x}}_{r},\boldsymbol{c}_{mv}\\}$,
where $\boldsymbol{c}_{mv}$ is the introduced canonical cameras in MVDream.
Then, the rest of diffusion loss is the same as Eqn.(1).

###  3.3 Multi-level Controllers

In order to insert the image prompt to control the output MV images, we
consider a multi-level strategy. The overall structure of the multi-level
controller from an image prompt can be seen in Fig. 4, and we elaborate the
details of each component in the following.

Global Controller. In our initial approach, we integrated global CLIP image
features into MVDream, akin to how text features are used, by fine-tuning the
model’s already well-established training. Recognizing that MVDream is
primarily trained on text embeddings, we introduced a multi-layer perceptron
(MLP) $\theta_{g}$, functioning as an adaptor similar to IP-Adaptor [51],
following the CLIP image global embedding. This step aims to align image
features with text features, ensuring compatibility within the MVDream
framework. Specifically, CLIP image encoding encodes image feature to a 1024
vector with a token length of 4, which we named as $\boldsymbol{f}_{g}$. And,
$\theta_{g}$ further adapts the image feature to be 1024 as the input to
cross-attention.

On the MVDiffusion side, inside of an attention layer $l$, we add a new set of
MLPs, $\theta_{k_{g},l}$ and $\theta_{v_{g},l}$, that takes the input the
adapted features and output its attention key and value matrix, which then
aggregated based on the query feature matrix, $\boldsymbol{q}_{l}$, yielding a
corresponding image cross-attention feature $\boldsymbol{h}_{g,l}$. Here, a
weight $\lambda=1.0$ is introduced to balance the hidden from text and image,
and the final output of layer $l$ is
$\boldsymbol{h}_{l}=\boldsymbol{h}_{t,l}+\lambda\boldsymbol{h}_{g,l}$. We
refer to decoupled cross-attention in IP-Adaptor for additional details.

To train such a model, we freeze the diffusion model, and only fine-tune
$\\{\theta_{g},\theta_{k_{g},l},\theta_{v_{g},l}\\}_{l}$. We follow the
training setting of MVDream by considering both 3D rendered datasets and 2D
image datasets together, which will be elaborated in our experimental section.

After the model is tuned, we found the model is able of absorb some
informations from the image such as structure of the object etc. As
illustrated in Fig. 4 (a), comparing with the input, the diffused output is
able to put the pirate hat similar to the image on the bulldog, while some
detailed pose and appearance information is lost, which we think is not enough
for a good control from the input image.

Local Controller. To enhance control, we try to utilize the hidden feature
from the CLIP encoder before its global pooling, which likely contains more
detailed structural information. This hidden feature, denoted as
$\boldsymbol{f}_{h}$, has a token length of $257$ and a feature dimension of
$1280$. A MLP adaptor $\theta_{h}$ is introduced to feed $\boldsymbol{f}_{h}$
into the diffusion network’s cross-attention module, with $\theta_{k_{h},l}$
and $\theta_{v_{h},l}$ forming the key and values matrix. These parameters,
$\\{\theta_{k_{h},l},\theta_{v_{h},l},\theta_{h}\\}_{l}$, are then jointly
trained as learnable elements similar to the global controller. Post-training,
we observed that the results were overly sensitive to image tokens, leading to
overexposed and unrealistic images, especially with higher class free guidance
(CFG) settings [33], as shown in Fig. 4(b).

To mitigate this, we implemented a resampling module $\theta_{r}$, following
the approach of IP-Adaptor, reducing the hidden token count from 257 to 16,
resulting in a more balanced local image feature $\boldsymbol{f}_{r}$. The
corresponding local controller parameters are
$\\{\theta_{r},\theta_{k_{r},l},\theta_{v_{r},l}\\}_{l}$. As Fig. 4(c)
illustrates, after this resampling, the diffused images more realistic, even
at higher CFG levels. From the generated images, it’s evident that the model
captures the overall layout and object shape, but also struggles with finer
identity details like object skin texture.

Figure 5: Example of artifacts we fixed with image-prompt score distillation
(see Sec. 3.4).

Pixel Controller. To optimally integrate object appearance texture, we propose
embedding the image prompt pixel latent ${\boldsymbol{x}}$ across all
attention layers in ImageDream. Specifically, MVDream employs a 3D dense self-
attention mechanism with a shape of $(bz,4,c,h_{l},w_{l})$ across four views
within a transformer layer. In contrast, ImageDream introduces an additional
frame by concatenating the input image, resulting in a feature shape of
$(bz,5,c,h_{l},w_{l})$. This enables similar 3D self-attention processes
between the four-view images and the input image.

During the training of our diffusion network, we refrain from adding noise to
the latent from the input image prompt, ensuring the network clearly captures
the image information. Additionally, to differentiate the input image features
and avoid confusion, we assign an all-zero vector to the camera embedding of
the input image. Given that the pixel controller is integrated into the multi-
view diffusion without extra parameters, we fine-tune all feature parameters
in unison, adopting the same training regime as the global/local controllers
but with a learning rate reduced by a factor of ten. This approach preserves
the original feature representations more effectively. Post-training, as
depicted in Fig 4(d), the generated multi-view images not only ethically
maintain the appearance from the input image but also uphold the multi-view
consistency characteristic of MVDream, resulting in satisfactory 3D model
fusion.

Finally, our multi-level controller is a combined one with local and pixel,
since we think the global one do not have too much additional information.
There might be potential queries regarding the necessity of a pixel
controller, given that IP-Adaptor, relying solely on CLIP features, can
capture extensive image texture details: we posit that while IP-Adaptor is
effective for modifying objects in the same view as the input, decoding the
same view is comparatively simpler. Multi-view diffusion, however, presents a
more complex challenge. Decoding from a highly compressed CLIP feature could
necessitate prolonged training on larger datasets. Therefore, at this stage,
we find the pixel controller significantly beneficial for rapidly training a
robust multi-view diffusion model.

###  3.4 Image-Prompt Sore Distillation

Implementing the image-prompt multi-view diffusion network in ImageDream
follows the multi-view score distillation framework of MVDream (Sec. 3.1),
with the addition of an image prompt as an input to the diffusion network.
However, we need to condier few key differences in NeRF optimization to
achieve accurate results.

Background Alignment. During SDS optimization, the NeRF-rendered image
includes a randomly colored background to differentiate the interior and
exterior of the 3D object. This random background, when input into the
diffusion network alongside the object, can conflict with the background from
the image prompt, leading to floating artifacts in the generated NeRF model,
as shown in Fig. 5(a). To resolve this, we adjusted the image-prompt
background to match the rendered background color from NeRF, successfully
eliminating these artifacts.

Camera Alignment. Our diffusion network tends to generate multi-view images
mirroring the camera parameters (e.g., elevation, field of view (FoV)) of the
input image prompt, parameters which remain unknown during NeRF rendering.
Randomly sampling parameters for rendering, as done in MVDream, can result in
images incongruent with the image prompt’s rendering settings, affecting the
geometry of detailed image structures. To mitigate this, we narrowed the
parameter sampling range from MVDream’s $[15,60]$, $[0,30]$ for camera FoV and
elevation to $[45,50]$ and $[0,5]$, respectively, a range more typical for a
generated user photos. This adjustment significantly improved the geometric
accuracy of the 3D objects, as demonstrated in Fig. 5(b).

We acknowledge this solution’s limitations; when the image prompt’s camera
parameters greatly differ from our selected range in canonical camera setting,
the resulting 3D object shape may be unpredictable. Future improvements could
include a camera parameter estimation module or increased randomness in the
image prompt rendering during diffusion training, to better synchronize the
settings between NeRF rendering and diffusion.

##  4 Experiments

Figure 6: User study of different methods. ImageDream-P: our full model w pixel controller, and ImageDream-G: without pixel controller (see Sec. 4.1).  | Synthesized Image | Re-rendered
---|---|---
Model | QIS(256)$\uparrow$ | CLIP(TX)$\uparrow$ | CLIP(IM)$\uparrow$ | QIS(320)$\uparrow$ | CLIP(TX)$\uparrow$ | CLIP(IM)$\uparrow$
SD-XL [1] | 52.0 ± 30.5 | 34.6 ± 3.09 | 100 | - | - | -
MVDream [37] | 23.05 ± 14.4 | 31.64 ± 2.99 | 78.41 ± 5.32 | 29.02 ± 10.24 | 32.69 ± 3.39 | 79.63 ± 4.15
Zero123 [19] | 22.16 ± 11.16 | 30.42 ± 3.19 | 84.88 ± 5.12 | - | - | -
Zero123-XL [19] | 34.07 ± 11.64 | 30.80 ± 2.59 | 84.10 ± 4.76 | 28.66 ± 5.03 | 29.19 ± 3.60 | 79.92 ± 6.59
Magic123 [31] | - | - | - | 24.23 ± 7.68 | 29.56 ±4.73 | 82.50 ± 8.78
SyncDreamer [20] | 22.04 ± 11.9 | 27.96 ± 3.01 | 78.17 ± 6.13 | 19.84 ± 6.64 | 25.82 ± 3.39 | 73.20 ± 6.30
ImageDream |  |  |  |  |  |
\- global | 22.31 ± 7.59 | 32.01 ± 2.84 | 84.50 ± 3.96 | 22.51 ± 5.86 | 31.48 ± 3.32 | 82.58 ± 4.35
\- local (-G) | 22.49 ± 9.57 | 31.32 ± 2.86 | 82.99 ± 6.03 | 22.30 ± 4.47 | 31.71 ± 2.96 | 84.34 ± 3.13
\- +pixel (-P) | 27.10 ± 12.8 | 32.39 ± 2.78 | 85.69 ± 3.77 | 25.16 ± 6.49 | 31.59 ± 3.23 | 84.83 ± 4.08
Table 1: Quantitative assessment of image synthesis quality using the test
prompt list from MVDream. The DDIM sampler was employed for testing, and
Implicit Volume from threestudio was utilized for re-rendering 3D model
images. Here Zero123 we took their checkpoint of 165K which including all
training instances. ’Zero123-XL’ indicates Zero123 trained with the larger
Objaverse-xl 10M dataset [7], while others were trained solely on the standard
Objaverse dataset [8], which is $\sim$10 times smaller. (-P) and (-G) are
correspondent to ImageDream-P and ImageDream-G in Fig.6. Figure 7: The
illustration of synthesized images from different baselines. Diffused: From
diffusion models. Reconstructed: Re-rendered images from correspondent fused
NeRF model. Please check the webpage for more results.

In this section, we detail the experimental setup for ImageDream, designed to
enable replication of our model. We will release both the model and code
following this submission.

Implementation Details. Adhering to the dataset configuration of MVDream (Sec.
3.3), we used a combined dataset from Objaverse for 3D multi-view rendering
and a 2D image dataset for training controllers (Sec. 3.3). For image prompts
in the 3D dataset, we randomly selected one of the 16 front-side views, with
azimuth angles ranging from $[-90,90]$ degrees, out of the total 32 circle
views. For the 2D dataset, we used the input image as the image prompt. A
random dropout rate of 0.1 was set for the image prompt during training,
replacing dropped prompts with a random uni-colored image. For all
experiments, i.e. with global controller, local controller and local plus
pixel controllers, we trained for $60$K steps with a batch size of 256 and a
gradient accumulation of 2, using the AdamW optimizer. The model is
initialized from stable diffusion (2.1) checkpoint of MVDream. The learning
rate was set to 1e-4, except for the model with pixel controller, where it was
reduced to 1e-5. Test image prompts were resized to $256\times 256$, and we
set the diffusion CFG to 5.0. The training takes $\sim$ 2 days with $8$ A100.
For NeRF optimization, we followed MVDream’s configuration but introduced a
three-stage optimization at resolutions $[64,192,256]$ which switched at
[$5$K, $10$K] steps, setting the camera distance between $[0.6,0.85]$ for
better NeRF model coverage. The NeRF training is about 1hr with A100.

Test Dataset. Our primary focus was evaluating ImageDream outside the
Objaverse distribution to ensure its practical applicability. We selected 39
well-curated prompts from MVDream, covering a diverse range of objects with
relatively complex geometries and appearances, surpassing datasets like
ShapeNet [5] or CO3D [18]. Using SDXL [29], we generated multiple images from
each prompt, selecting ones with aesthetically pleasing objects. The
backgrounds of these images were then removed, and the objects re-centered,
akin to the approach used in Zero123.

###  4.1 Comparisons

In our evaluation, we compared ImageDream’s performance against several SoTA
baselines, including Zero123-XL [19] (trained on 10x larger data than ours),
Magic123 [31], and SyncDreamer [20]. The criteria for comparison were geometry
quality and similarity to the image prompt (IP). ’Geometry quality’ refers to
the generated 3D asset’s conformance to common sense in terms of shape and
minimal artifacts, while ’similarity to IP’ assesses the resemblance of the
results to the input image. We executed all baseline tests using default
configurations as implemented in threestudio111https://github.com/threestudio-
project/threestudio.

Qualitative Evaluation. Lacking ground truth for the test image prompts, we
conducted a real user study to evaluate the quality of the generated 3D
models. Participants were briefed on our evaluation standards and asked to
choose their preferred model based on these criteria. The experiment was
double-blind, with participants shown 3D assets generated by different methods
without identifying labels. The comparison results, depicted in Fig. 6, show
that ImageDream, both with (ImageDream-P) and without (ImageDream-G) the pixel
controller, significantly outperformed other baselines. ImageDream-P was
particularly favored, while ImageDream-G also received a positive preference
rate. SyncDreamer was omitted from the figure due to its NeuS results having a
0$\%$ preference rate.

Fig. 7 presents a representative case comparing results from the diffusion
models and the final NeRF model. Systems like Magic123 and Zero123, which rely
on single-view diffusion with relative camera embedding, often produce
incorrect geometry, as illustrated by their inability to accurately represent
the span of horse body. In contrast, ImageDream, through its unique design,
effectively resolves this issue, resulting in more satisfactory models (more
results are list in webpage).

Numerical Evaluation. To thoroughly assess image quality at various stages of
our generation pipeline, we employed the Inception Score (IS)[34] and CLIP
scores[32] using text-prompt and image-prompt, respectively. The IS evaluates
image quality, while CLIP scores assess text-image and image-image alignment.
However, since IS traditionally evaluates both image quality and diversity
within a set, and our prompt quantity is limited, the diversity aspect makes
the score less reliable. Therefore, we modified the IS by omitting its
diversity evaluation, replacing the mean distribution with a uniform
distribution. Specifically, we set $q_{i}$ in IS to be $1/N$, making the IS of
an image $\sum_{i}p_{i}\log(Np_{i})$, where $N$ is the inception class count
and $p_{i}$ is the predicted probability for the $i_{th}$ class. We denote
this modified metric as Quality-only IS (QIS). For the CLIP score, we
calculated the mean score between each generated view and the provided text-
prompt or image-prompt.

In Tab.1, we present comparative results. SD-XL, reflecting the score of test
images, achieved the highest QIS and CLIP scores. MVDream, listed as a
benchmark for final 3D model quality, shows improved synthesized image quality
after 3D fusion due to multi-view consistency. In contrast, Zero123 and
Zero123-XL experienced a drop in image quality post-3D fusion due to diffusion
inconsistency. Magic123 enhanced the CLIP score over Zero123 by integrating a
joint diffusion model. SyncDreamer’s quality declined as it diffuses only 16
fixed views, complicating reconstruction. In ImageDream, we evaluated three
models for ablation: one with a global controller, another with a local
controller, and the last incorporating both local and pixel controllers
(Sec.3.3). ImageDream maintained high image quality in both diffusion and
post-3D fusion stages. The local controller, in particular, provided better
image CLIP scores post-fusion, thanks to richer image feature representations.
The pixel controller model excelled in image CLIP scores during both stages.
Notably, ImageDream-pixel ranked second in other scores, with Zero123-XL using
a significantly larger dataset (Objaverse-XL [7]).

However, these scores don’t fully encapsulate important aspects like multi-
view consistency and geometric correctness. For instance, as Fig. 7
demonstrates, zero123-XL, despite having high IS due to easy image
classification, showed poorer consistency. Thus, while these scores offer some
reliability when consistency is high, future research should aim to develop
more comprehensive metrics that accurately capture geometric correctness to
better compare different generation algorithms.

Figure 8: Illustration of a failure case. Where Trumps’ face turns blurry
after 3D fusion reconstruction (first row). However, the face textures can be
generated with model using global-only controller thanks to its semantic
global representation.

###  4.2 Limitations

While the model incorporating the pixel controller achieves the best scores,
we observed certain trade-offs, particularly when the image constraints are
overly stringent. For instance, in cases like the small facial details of a
full-body avatar (shown in Fig.8), the model struggles to capture these
nuances, whereas global control might recover the face based on the text
prompt. To address this, as outlined in Sec.3.4, the pixel controller model
needs to better estimate image intrinsic and extrinsic properties, or a better
balance tuning inside of multi-level controllers. This may be solved by
exploring the use of larger models, such as SDXL [29], which be our future
work.

##  5 Conclusion

We introduce ImageDream, an advanced image-prompt 3D generation model
utilizing multi-view diffusion. This model innovatively applies canonical
camera coordination and multi-level image-prompt controllers, enhancing
control and addressing geometric inaccuracies seen in prior methods. Future
improvements could focus on increasing randomness in image-prompts during
training to further reduce texture blurriness in the generated models. These
steps are expected to further advance the capabilities and applications of
ImageDream in 3D model generation.

##  6 Acknowledgements and Ethics Statement

We thank our 3D group members of Kejie Li, and our intern Zeyuan Chen in joint
meeting discussion and setup baselines of SyncDreamer, which help complete
this paper. In addition, note that the models proposed in this paper aims to
facilitate the 3D generation task that is widely demanded in industry for
ethical purpose. It could be potentially applied to unwanted scenarios such as
generating violent and sexual content by third-party fine-tuning. Built upon
the Stable Diffusion model [33], it might also inherit the biases and
limitations to generate unwanted results. Therefore, we believe that the
images or models synthesized using our approach should be carefully examined
and be presented as synthetic. Such generative models may also have the
potential to displace creative workers via automation. That being said, these
tools may also enable growth and improve accessibility for the creative
industry.

## References

  * [1] stable-diffusion-xl-base-1.0.  https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0.  Accessed: 2023-08-29.
  * [2] Stable diffusion image variation.  https://huggingface.co/spaces/lambdalabs/stable-diffusion-image-variations.
  * Chan et al. [2022] Eric R Chan, Connor Z Lin, Matthew A Chan, Koki Nagano, Boxiao Pan, Shalini De Mello, Orazio Gallo, Leonidas J Guibas, Jonathan Tremblay, Sameh Khamis, et al.  Efficient geometry-aware 3d generative adversarial networks.  In _CVPR_ , 2022.
  * Chan et al. [2023] Eric R. Chan, Koki Nagano, Matthew A. Chan, Alexander W. Bergman, Jeong Joon Park, Axel Levy, Miika Aittala, Shalini De Mello, Tero Karras, and Gordon Wetzstein.  GeNVS: Generative novel view synthesis with 3D-aware diffusion models.  In _arXiv_ , 2023.
  * Chang et al. [2015] Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al.  Shapenet: An information-rich 3d model repository.  _arXiv preprint arXiv:1512.03012_ , 2015.
  * Chen et al. [2023] Rui Chen, Yongwei Chen, Ningxin Jiao, and Kui Jia.  Fantasia3d: Disentangling geometry and appearance for high-quality text-to-3d content creation.  _arXiv:2303.13873_ , 2023.
  * Deitke et al. [2023a] Matt Deitke, Ruoshi Liu, Matthew Wallingford, Huong Ngo, Oscar Michel, Aditya Kusupati, Alan Fan, Christian Laforte, Vikram Voleti, Samir Yitzhak Gadre, Eli VanderBilt, Aniruddha Kembhavi, Carl Vondrick, Georgia Gkioxari, Kiana Ehsani, Ludwig Schmidt, and Ali Farhadi.  Objaverse-xl: A universe of 10m+ 3d objects.  2023a.
  * Deitke et al. [2023b] Matt Deitke, Dustin Schwenk, Jordi Salvador, Luca Weihs, Oscar Michel, Eli VanderBilt, Ludwig Schmidt, Kiana Ehsani, Aniruddha Kembhavi, and Ali Farhadi.  Objaverse: A universe of annotated 3d objects.  In _CVPR_ , pages 13142–13153, 2023b.
  * Deng et al. [2022] Yu Deng, Jiaolong Yang, Jianfeng Xiang, and Xin Tong.  Gram: Generative radiance manifolds for 3d-aware image generation.  In _CVPR_ , pages 10673–10683, 2022.
  * Gao et al. [2022] Jun Gao, Tianchang Shen, Zian Wang, Wenzheng Chen, Kangxue Yin, Daiqing Li, Or Litany, Zan Gojcic, and Sanja Fidler.  Get3d: A generative model of high quality 3d textured shapes learned from images.  _NeurIPS_ , 2022.
  * Henderson and Ferrari [2020] Paul Henderson and Vittorio Ferrari.  Learning single-image 3d reconstruction by generative modelling of shape, pose and shading.  _International Journal of Computer Vision_ , 2020.
  * Henderson et al. [2020] Paul Henderson, Vagia Tsiminaki, and Christoph H Lampert.  Leveraging 2d data to learn textured 3d mesh generation.  In _CVPR_ , 2020.
  * Huang et al. [2023] Yukun Huang, Jianan Wang, Yukai Shi, Xianbiao Qi, Zheng-Jun Zha, and Lei Zhang.  Dreamtime: An improved optimization strategy for text-to-3d content creation.  _arXiv:2306.12422_ , 2023.
  * Jun and Nichol [2023] Heewoo Jun and Alex Nichol.  Shap-e: Generating conditional 3d implicit functions.  _arXiv:2305.02463_ , 2023.
  * Karnewar et al. [2023] Animesh Karnewar, Andrea Vedaldi, David Novotny, and Niloy J Mitra.  Holodiffusion: Training a 3d diffusion model using 2d images.  In _CVPR_ , 2023.
  * Kingma and Welling [2014] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  In _ICLR_ , 2014.
  * Lin et al. [2023a] Chen-Hsuan Lin, Jun Gao, Luming Tang, Towaki Takikawa, Xiaohui Zeng, Xun Huang, Karsten Kreis, Sanja Fidler, Ming-Yu Liu, and Tsung-Yi Lin.  Magic3d: High-resolution text-to-3d content creation.  In _CVPR_ , 2023a.
  * Lin et al. [2023b] Shanchuan Lin, Bingchen Liu, Jiashi Li, and Xiao Yang.  Common diffusion noise schedules and sample steps are flawed.  _arXiv:2305.08891_ , 2023b.
  * Liu et al. [2023a] Ruoshi Liu, Rundi Wu, Basile Van Hoorick, Pavel Tokmakov, Sergey Zakharov, and Carl Vondrick.  Zero-1-to-3: Zero-shot one image to 3d object.  _arXiv:2303.11328_ , 2023a.
  * Liu et al. [2023b] Yuan Liu, Cheng Lin, Zijiao Zeng, Xiaoxiao Long, Lingjie Liu, Taku Komura, and Wenping Wang.  Syncdreamer: Learning to generate multiview-consistent images from a single-view image.  _arXiv preprint arXiv:2309.03453_ , 2023b.
  * Long et al. [2023] Xiaoxiao Long, Yuan-Chen Guo, Cheng Lin, Yuan Liu, Zhiyang Dou, Lingjie Liu, Yuexin Ma, Song-Hai Zhang, Marc Habermann, Christian Theobalt, and Wenping Wang.  Wonder3d: Single image to 3d using cross-domain diffusion, 2023.
  * Melas-Kyriazi et al. [2023] Luke Melas-Kyriazi, Iro Laina, Christian Rupprecht, and Andrea Vedaldi.  Realfusion: 360deg reconstruction of any object from a single image.  In _CVPR_ , 2023.
  * Mildenhall et al. [2021] Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng.  Nerf: Representing scenes as neural radiance fields for view synthesis.  In _ECCV_ , 2021.
  * Nguyen-Phuoc et al. [2019] Thu Nguyen-Phuoc, Chuan Li, Lucas Theis, Christian Richardt, and Yong-Liang Yang.  Hologan: Unsupervised learning of 3d representations from natural images.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2019.
  * Nguyen-Phuoc et al. [2020] Thu H Nguyen-Phuoc, Christian Richardt, Long Mai, Yongliang Yang, and Niloy Mitra.  Blockgan: Learning 3d object-aware scene representations from unlabelled images.  _NeurIPS_ , 2020.
  * Nichol et al. [2022] Alex Nichol, Heewoo Jun, Prafulla Dhariwal, Pamela Mishkin, and Mark Chen.  Point-e: A system for generating 3d point clouds from complex prompts.  _arXiv:2212.08751_ , 2022.
  * Niemeyer and Geiger [2021] Michael Niemeyer and Andreas Geiger.  Giraffe: Representing scenes as compositional generative neural feature fields.  In _CVPR_ , 2021.
  * Niemeyer et al. [2022] Michael Niemeyer, Jonathan T. Barron, Ben Mildenhall, Mehdi S. M. Sajjadi, Andreas Geiger, and Noha Radwan.  Regnerf: Regularizing neural radiance fields for view synthesis from sparse inputs.  In _CVPR_ , 2022.
  * Podell et al. [2023] Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach.  Sdxl: Improving latent diffusion models for high-resolution image synthesis.  _arXiv:2307.01952_ , 2023.
  * Poole et al. [2023] Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall.  Dreamfusion: Text-to-3d using 2d diffusion.  In _ICLR_ , 2023.
  * Qian et al. [2023] Guocheng Qian, Jinjie Mai, Abdullah Hamdi, Jian Ren, Aliaksandr Siarohin, Bing Li, Hsin-Ying Lee, Ivan Skorokhodov, Peter Wonka, Sergey Tulyakov, and Bernard Ghanem.  Magic123: One image to high-quality 3d object generation using both 2d and 3d diffusion priors.  _arXiv preprint arXiv:2306.17843_ , 2023.
  * Radford et al. [2021] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _ICML_ , 2021.
  * Rombach et al. [2022] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _CVPR_ , 2022.
  * Salimans et al. [2016] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen.  Improved techniques for training gans.  _NeurIPS_ , 2016.
  * Schuhmann et al. [2022] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al.  Laion-5b: An open large-scale dataset for training next generation image-text models.  _NeurIPS_ , 2022.
  * Shi et al. [2023a] Ruoxi Shi, Hansheng Chen, Zhuoyang Zhang, Minghua Liu, Chao Xu, Xinyue Wei, Linghao Chen, Chong Zeng, and Hao Su.  Zero123++: a single image to consistent multi-view diffusion base model.  _arXiv preprint arXiv:2310.15110_ , 2023a.
  * Shi et al. [2023b] Yichun Shi, Peng Wang, Jianglong Ye, Long Mai, Kejie Li, and Xiao Yang.  Mvdream: Multi-view diffusion for 3d generation.  _arXiv:2308.16512_ , 2023b.
  * Shue et al. [2023] J Ryan Shue, Eric Ryan Chan, Ryan Po, Zachary Ankner, Jiajun Wu, and Gordon Wetzstein.  3d neural field generation using triplane diffusion.  In _CVPR_ , 2023.
  * Sitzmann et al. [2019] Vincent Sitzmann, Michael Zollhöfer, and Gordon Wetzstein.  Scene representation networks: Continuous 3d-structure-aware neural scene representations.  _NeurIPS_ , 32, 2019.
  * Sun et al. [2023] Jingxiang Sun, Bo Zhang, Ruizhi Shao, Lizhen Wang, Wen Liu, Zhenda Xie, and Yebin Liu.  Dreamcraft3d: Hierarchical 3d generation with bootstrapped diffusion prior, 2023.
  * Szymanowicz et al. [2023] Stanislaw Szymanowicz, Christian Rupprecht, and Andrea Vedaldi.  Viewset diffusion:(0-) image-conditioned 3d generative models from 2d data, 2023.
  * Tang et al. [2023] Junshu Tang, Tengfei Wang, Bo Zhang, Ting Zhang, Ran Yi, Lizhuang Ma, and Dong Chen.  Make-it-3d: High-fidelity 3d creation from a single image with diffusion prior.  _arXiv:2303.14184_ , 2023.
  * Tsalicoglou et al. [2023] Christina Tsalicoglou, Fabian Manhardt, Alessio Tonioni, Michael Niemeyer, and Federico Tombari.  Textmesh: Generation of realistic 3d meshes from text prompts.  _arXiv:2304.12439_ , 2023.
  * Wang et al. [2023a] Haochen Wang, Xiaodan Du, Jiahao Li, Raymond A Yeh, and Greg Shakhnarovich.  Score jacobian chaining: Lifting pretrained 2d diffusion models for 3d generation.  In _CVPR_ , 2023a.
  * Wang et al. [2023b] Tengfei Wang, Bo Zhang, Ting Zhang, Shuyang Gu, Jianmin Bao, Tadas Baltrusaitis, Jingjing Shen, Dong Chen, Fang Wen, Qifeng Chen, et al.  Rodin: A generative model for sculpting 3d digital avatars using diffusion.  In _CVPR_ , 2023b.
  * Wang et al. [2023c] Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu.  Prolificdreamer: High-fidelity and diverse text-to-3d generation with variational score distillation.  _arXiv:2305.16213_ , 2023c.
  * Watson et al. [2023] Daniel Watson, William Chan, Ricardo Martin-Brualla, Jonathan Ho, Andrea Tagliasacchi, and Mohammad Norouzi.  Novel view synthesis with diffusion models.  In _ICLR_ , 2023.
  * Wu et al. [2023] Chao-Yuan Wu, Justin Johnson, Jitendra Malik, Christoph Feichtenhofer, and Georgia Gkioxari.  Multiview compressive coding for 3d reconstruction.  In _CVPR_ , 2023.
  * Xu et al. [2022a] Dejia Xu, Yifan Jiang, Peihao Wang, Zhiwen Fan, Humphrey Shi, and Zhangyang Wang.  Sinnerf: Training neural radiance fields on complex scenes from a single image.  2022a.
  * Xu et al. [2022b] Dejia Xu, Yifan Jiang, Peihao Wang, Zhiwen Fan, Yi Wang, and Zhangyang Wang.  Neurallift-360: Lifting an in-the-wild 2d photo to a 3d object with 360° views.  2022b.
  * Ye et al. [2023a] Hu Ye, Jun Zhang, Sibo Liu, Xiao Han, and Wei Yang.  Ip-adapter: Text compatible image prompt adapter for text-to-image diffusion models.  _arXiv preprint arXiv:2308.06721_ , 2023a.
  * Ye et al. [2023b] Jianglong Ye, Peng Wang, Kejie Li, Yichun Shi, and Heng Wang.  Consistent-1-to-3: Consistent image to 3d view synthesis via geometry-aware diffusion models.  _arXiv preprint arXiv:2310.03020_ , 2023b.
  * Zhou and Tulsiani [2023] Zhizhuo Zhou and Shubham Tulsiani.  Sparsefusion: Distilling view-conditioned diffusion for 3d reconstruction.  In _CVPR_ , 2023.


# Conformal Inference for Invariant Representation

Wenlu TANG
Department of Applied Mathemetics
Hong Kong Polytechnic University
<EMAIL_ADDRESS>
&Zicheng LIU
Department of Statistics
The Chinese University of Hong Kong
<EMAIL_ADDRESS>

###### Abstract

The application of machine learning models can be significantly impeded by the
occurrence of distributional shifts, as the assumption of homogeneity between
the population of training and testing samples in machine learning and
statistics may not be feasible in practical situations. One way to tackle this
problem is to use invariant learning, such as invariant risk minimization
(IRM), to acquire an invariant representation that aids in generalization with
distributional shifts. This paper develops methods for obtaining distribution-
free prediction regions to describe uncertainty estimates for invariant
representations, accounting for the distribution shifts of data from different
environments. Our approach involves a weighted conformity score that adapts to
the specific environment in which the test sample is situated. We construct an
adaptive conformal interval using the weighted conformity score and prove its
conditional average under certain conditions. To demonstrate the effectiveness
of our approach, we conduct several numerical experiments, including
simulation studies and a practical example using real-world data.

##  1 Introduction

The assumption of independent and identically distributed training and test
data is a fundamental principle of statistical machine learning. However, in
real-world scenarios, distributional shifts are common and can pose challenges
for machine learning models when applied to new data (Sugiyama et al., (2007),
Taori et al., (2020)). To address these challenges, researchers have proposed
various strategies to develop predictive models that can adapt to changes in
the distribution of the data. The negative impact of such shifts on the
quality of predictions can be significant, and it is therefore essential to
construct prediction intervals that not only provide point prediction but also
quantify the uncertainty of predictions under distributional shifts.

Conformal inference has become a powerful and valuable tool for constructing
prediction intervals. Conformal prediction is a distribution-free approach
that constructs prediction intervals with a guaranteed coverage probability in
finite samples. The objective is to create a prediction set $C_{n}$ that
includes $Y_{n+1}$ with a probability of at least $1-\alpha$ (Vovk et al.,
(2005)). Since the pioneering work by Vovk et al., (2005), there have been
numerous follow-up studies and extensions in both computation and theory. Lei
et al., (2013), Lei and Wasserman, (2014), and Lei et al., (2018) have
developed much insightful work on statistical theory for conformal methods.
Marginal validity, a conventional coverage guarantee that can be achieved
under the i.i.d assumption, is demonstrated in Vovk et al., (2005) and Zeni et
al., (2020). However, as demonstrated in Lei and Wasserman, (2014) and Vovk,
(2012), conditional validity with a finite-length prediction interval is
impossible without regularity and consistency assumptions on the model and
estimator. Recently, Romano et al., (2019) and Kivaranovic et al., (2020)
proposed a technique that combines conformal prediction with classical
quantile regression. Lei et al., (2018) prove that conformal prediction
intervals are accurate under minimal assumptions on the residuals, indicating
that they do not substantially over-cover. Besides the theoretical results,
the computational efficiency of the conformal method is also attractive. The
mentioned conformal methods rely on the assumption of sample exchangeability,
but distribution shifts are prevalent, which can violate this assumption.

Recently, numerous studies in conformal inference have aimed to extend the
method to better adapt to distribution shifts. Gibbs and Candes,
(2021)constructed prediction sets in an online setting where the data
generating distribution can change over time in an unknown way. Meanwhile,
Zaffran et al., (2022) analyzed the impact of this method in time series.
Tibshirani et al., (2019) proposed a method to obtain accurate prediction sets
for covariate-shifted data by weighting calibration data using the likelihood
ratio, while Barber et al., (2022) used a non-symmetric algorithm that assigns
more weight to recent observations. Cauchois et al., (2020) examined
distribution shift in an $f$-divergence ball and proposed using distributional
robust optimization to create prediction sets. Guan, (2023) expanded the
conformal prediction framework by offering a locally-focused, single-test-
sample adaptive construction that can be adapted to various conformal scores.
Dunn et al., (2022) extended conformal methods to a hierarchical setting by
using data from multiple training domains.

This paper aims to quantify the uncertainty estimates of distribution shift
through out-of-distribution generalization via Invariant Risk Minimization.
This approach assumes that the variation across the training domain represents
the variation in the testing domain.

###  1.1 Invariant Risk Minimization

Even though neural networks can perform exceptionally well on the training
data by Empirical Risk Minimization (ERM), where the average loss is
considered across all training environments, they can be highly sensitive to
distributional shifts, which makes it challenging to apply them in practical
situations. Therefore, it is essential to focus on identifying and addressing
distributional shifts to improve the performance of machine learning models in
real-world applications. Invariant risk minimization (IRM) (Arjovsky et al.,
(2019)), is one of the approaches aims to learn models that are robust to
distributional shifts in the data.

Let us suppose data are coming from a collection of $m$ environments denoted
by $\mathcal{E}_{tr}$, which includes environments $e\in\\{1,\ldots,m\\}$,
each defined on the sample space $\mathcal{X}\times\mathcal{Y}$. Here,
$\mathcal{X}$ and $\mathcal{Y}$ denote the input and target spaces,
respectively. The joint distribution in distinct environments is denoted by
$\operatorname{Pr}^{e}(\boldsymbol{x},\boldsymbol{y})$, where
$e\in\mathcal{E}$. Let
$\mathcal{D}^{e}:=\left\\{\left(\mathbf{x}_{i}^{e},\mathbf{y}_{i}^{e}\right)\right\\}_{i=1}^{n_{e}}$
be the data set drawn from $e\in\mathcal{E}_{tr}$ with $n_{e}$ being data set
size. Suppose that data
${\mathcal{D}}^{e}=(x_{i}^{e},y_{i}^{e})_{i=1}^{n_{e}}$ is collected from
different environments $e\in\\{1,\ldots,m\\},$ and $n_{e}$ is the sample size
in each environment. The basic idea behind IRM is to identify the underlying
causal mechanisms without environment spurious effect and to explicitly model
them. This allows the model to generalize across different environments, even
if they have different distributions. Essentially, the goal is to learn a
model that is invariant to changes in the data distribution. The main
invariance assumption in IRM is the existence of a data representation
$\Phi(X)$ such that for all $e,e^{\prime}\in\mathcal{E}_{tr}$,
$\mathbb{E}\left[Y\mid\Phi\left(X^{e}\right)\right]=\mathbb{E}\left[Y\mid\Phi\left(X^{e^{\prime}}\right)\right]$,
where $\mathcal{E}_{{tr}}$ denotes the available training environments. To
achieve this goal, IRM uses a regularized optimization approach that
encourages the model to learn feature representation $\Phi(x)$ and optimal
classifier $\omega$ that are predictive of the target variable, but are also
invariant to changes in the data distribution. Existing IRM methods learn
$\omega$ and $\Phi(x)$ by solving the following optimization problem:

$\displaystyle\min_{\omega,\Phi}\sum_{e\in\mathcal{E}_{tr}}\mathcal{R}^{e}(\omega,\Phi)+\lambda\mathcal{J}(\omega,\Phi),$

where
$\mathcal{R}^{e}(\omega,\Phi)={\mathbb{E}}_{X^{e},Y^{e}}[l(\omega\cdot\Phi(X^{e}),Y^{e})]$
is the risk under environment $e$ and $l$ is the loss function. The
regularizer penalty $\mathcal{J}(\omega,\Phi)$ is designed to ensure that the
model learns features that are common across all environments and it varies in
different approaches. In Arjovsky et al., (2019), they propose IRMv1 where
$\mathcal{J}(\omega,\Phi)=\left\|\nabla_{\omega}R^{e}(\omega\cdot\Phi)\right\|^{2}$.
Several studies have suggested variations to this objective by introducing
alternative regularizer to fit stronger invariance assumptions. Under a
stronger invariance assumption that
$\mathbb{P}[Y\mid\Phi(X^{e})]=\mathbb{P}[Y\mid\Phi(X^{e^{\prime}})]$, Krueger
et al., (2021) proposed REx with
$\mathcal{J}(\omega,\Phi)=\operatorname{Var}\left(\mathcal{R}^{e}(\omega,\Phi)\right)$
and Chang et al., (2020) proposed InvRat with
$\mathcal{J}(\omega,\Phi)=\lambda\left(R_{e}\left(\omega,\Phi\right)-R_{e}\left(\omega_{e},\Phi\right)\right)$.
Several follow-up works has been proposed to make the IRM more effective. Zhou
et al., (2022) add a sparsity constraint to the network and train a neural
network that is sparse to prevent overfitting. Lin et al., (2022) updated
InvRat by Bayesian method with a posterior distribution of the classifier.
Rosenfeld et al., (2020) points out some limitations with IRM on
classification tasks. Mahajan et al., (2021) introduces a novel regularizer to
match the representation of the same object in different environments. Wang et
al., (2022) proposes a simple post-processing method for solving the IRM
problem without retraining the model. Ahuja et al., (2020) searches for Nash
equilibrium solutions between several environments in a game and let all the
environments play the game to determine the action to minimize overall risk.
Chang et al., (2020), Koyama and Yamaguchi, (2020),Li et al., (2022) and Ahuja
et al., (2021) consider the invariant learning problem from the information
theory aspect. Creager et al., (2021) proposed Environment Inference for
Invariant Learning(EIIL) that attempts to automatically partition a dataset
into different environments to learn environments labels that maximize the
IRM’s penalty. IRM has been applied to a variety of machine learning tasks,
including classification, regression, and reinforcement learning. It has shown
promising results in settings where there are known or unknown distributional
shifts in the data, such as in healthcare, finance, and other domains where
data is constantly changing.

The methods of IRM vary in terms of objective functions, penalties, training
frameworks, and optimization techniques. Although IRM’s effectiveness is often
evaluated by comparing its classification accuracy on out-of-distribution
testing data to that of ERM, there is currently no specific criterion to
assess the performance of the invariant representation. Furthermore, there is
limited literature on the development of valid prediction intervals for
invariant learning problems. These gaps in the literature motivate the present
study.

###  1.2 Summary and Outline

In this paper, we make the following methodological and theoretical
contributions.

  * •

We propose a criterion to evaluate the invariance performance of the invariant
representation obtained in Invariant Risk Minimization (IRM). We demonstrate
that the expectation of an ideal invariant predictor across different
environments only varies due to covariate shift.

  * •

We perform split conformal inference on training data from various
environments by combining the data from each environment to construct a
conformal interval. We then evaluate the marginal coverage for this interval.

  * •

We propose a weighted conformity score, where the weight of a data point
$x_{i}$ corresponding to environment $e$ is determined by the similarity of
$x_{i}$ and samples in environment $e$. We then construct an adaptive
conformal interval using this weighted conformity score, which holds locally
coverage in each environment.

  * •

We establish theoretical guarantees for the valid coverage of our proposed
approach and provide a method to check the condition using the invariance
assessment criterion. Extensive numerical studies are conducted to support the
theory.

Table 1: A comparison of some recent conformal interval under distribution shift. Literature | Method | Multi-environment | Online | Conditional validity
---|---|---|---|---
Dunn et al., (2022) | Two-Layer | ✔ | ✗ | ✗
Gibbs and Candes, (2021) | ACI | ✗ | ✔ | ✗
Cauchois et al., (2020) | DRO | ✗ | ✗ | ✗
This paper | ACIR | ✔ | ✔ | ✔

In Section 2.1, we present a novel statistic for the invariant representation
that can assess the invariance of the representation. Section 2.2 introduces
the split conformal method for data from different environments, while Section
2.3 proposes a weighted conformity score based on the similarity between the
testing point and each environment. This weighted conformity score is adaptive
to the variation in different environments, and we use it to construct an
adaptive conformal prediction interval for the invariant predictions, which we
provide theoretical guarantees for. In Section 3, we conduct extensive
simulation studies to examine the finite-sample performance of the proposed
method. Finally, we provide concluding remarks in Section 4, and all technical
proofs are deferred to the Supplementary Materials.

##  2 Methodology

In this section, we propose a split conformal and adaptive conformal interval
for IRM predictions.

#### Preliminaries

Suppose that training data
${\mathcal{D}}^{e}=(x_{i}^{e},y_{i}^{e})_{i=1}^{n_{e}},$ is collected from
different environments $e=1,\ldots,m$ with joint distribution
$\operatorname{Pr}^{e}(x,y)$ in the environment $e$, where $n_{e}$ is the
sample size in each data set ${\mathcal{D}}^{e}$. Then IRMv1 (Arjovsky et al.,
(2019)) solves

$\min_{\Phi,\omega}\sum_{e\in\mathcal{E}_{\mathrm{tr}}}{\mathcal{R}}^{e}(\Phi,\omega)+\lambda\left\|\nabla_{w\mid
w=1.0}R^{e}(w,\Phi)\right\|^{2},$ (2.1)

where ${\mathcal{R}}^{e}(\Phi,\omega)$ is the risk w.r.t the environment
$e$,$\Phi:{\mathcal{X}}\to\mathcal{H}$ is the data representation mapping
covariates from covariate domain ${\mathcal{X}}$ to the embedded space
$\mathcal{H}$, and $\omega:\mathcal{H}\to\mathcal{Y}$ is the classifier in
classification or the last layer in multilayer perception in regression
problems. We can obtain the estimate for the invariant predictor
${\mathbb{E}}(Y|X=x)$, which is $\hat{f}(x)=\hat{\omega}\cdot\hat{\Phi}(x)$.

###  2.1 Invariance Assessment

In Arjovsky et al., (2019), the key idea of the invariance assumption is that
the expected value of the outcome variable given the invariant representation
$\Phi(\cdot)$ is identical across environments. Specifically, it assumes that
there exists invariant representation $\Phi(\cdot)$, such that for distinct
$e,e^{\prime}\in\mathcal{E}_{tr}$,
$\mathbb{E}\left[Y\mid\Phi\left(X^{e}\right)\right]=\mathbb{E}[Y\mid\Phi(X^{e^{\prime}})]$,
where $X^{e}$ is the covariate collected from environment $e$. If we take the
expectation w.r.t. $X^{e}$ on both sides, we get:

$\displaystyle{\mathbb{E}}_{X^{e}}\left\\{{\mathbb{E}}\left[Y\mid\Phi\left(X^{e}\right)\right]\right\\}=\int_{{\mathcal{X}}}{\mathbb{E}}\left[Y\mid\Phi\left(X^{e}\right)\right]d{\mathbb{P}}(X^{e})$
$\displaystyle=$
$\displaystyle\int_{{\mathcal{X}}}\mathbb{E}[Y\mid\Phi(X^{e^{\prime}})]d{\mathbb{P}}(X^{e})=\int_{{\mathcal{X}}}\mathbb{E}[Y\mid\Phi(X^{e^{\prime}})]d{\mathbb{P}}(X^{{e^{\prime}}})\frac{d{\mathbb{P}}(X^{{e}})}{d{\mathbb{P}}(X^{{e^{\prime}}})}$
$\displaystyle=$
$\displaystyle{\mathbb{E}}_{X^{e^{\prime}}}\left\\{\mathbb{E}[Y\mid\Phi(X^{e^{\prime}})]\rho(X^{{e^{\prime}}},X^{e})\right\\},$

where $\rho(X^{e^{\prime}},X^{e})$ is the the likelihood ratio of
$d{\mathbb{P}}(X^{{e}})/d{\mathbb{P}}(X^{{e^{\prime}}})$. The invariance
assumption posits that the expected value of the outcome variable given the
invariant representation $\Phi(\cdot)$ is the same across different
environments. This implies that the transformation of the mean value of the
invariant representation $w\cdot\Phi(\cdot)$ across environments only depends
on the covariate shift $\rho(X^{e^{\prime}},X^{e})$. Previous studies have
investigated the covariate shift $\rho(X^{{e^{\prime}}},X^{e})$, see Sugiyama
and Müller, (2005), Quinonero-Candela et al., (2008), Chen et al., (2016). If
we know or can accurately estimate the likelihood ratio between every pair of
environments, we can evaluate the performance of the invariant representation
$\Phi(\cdot)$. For instance, we can choose a baseline environment $e$ and
compute

$M_{e}(\Phi,X^{e^{\prime}})={\mathbb{E}}_{X^{e^{\prime}}}\left[\mathbb{E}[Y\mid{\Phi}(X^{e^{\prime}})]\rho(X^{{e^{\prime}}},X^{e})\right]$

for $e^{\prime}=1,\ldots,m$ and $\rho(X^{{e}},X^{e})=1$. It can be estimated
by empirical expectation

$\hat{M}_{e}(\Phi,X^{e^{\prime}})=\sum_{x_{i}:i\in
e^{\prime}}\hat{\omega}\cdot\hat{\Phi}(x_{i})\rho(X^{{e^{\prime}}},X^{e}).$

The variance of $M_{e}(\hat{\Phi},X^{e^{\prime}})$ can serve as a benchmark
for the invariance of the estimated invariant representation
$\hat{\Phi}(\cdot)$. Denote

${\text{Inv}}(\Phi)={\mathbb{E}}_{e}\\{Var_{e^{\prime}}(M_{e}(\Phi^{or},X^{e^{\prime}}))\\}$
(2.2)

as the average of the variance across the environment. For the optimal
invariant representation $\Phi^{or}(\cdot)$ that satisfies invariance
assumption, ${\text{Inv}}(\Phi^{or})=0$. If the estimate from the IRM model
approximates the optimal invariant representation $\Phi^{or}$, the value
${\text{Inv}}(\Phi)$ will be small. Table 2 shows the value of
${\text{Inv}}(\Phi)$ in the SEM setting same as the setting in Arjovsky et
al., (2019). It shows that the ${\text{Inv}}(\Phi)$ value in IRM model is
smaller than the ERM model.

Table 2: The invariance evaluation by ${\text{Inv}}(\hat{\Phi})$ value. Method | FOU | POU | FEU | PEU
---|---|---|---|---
ERM | 0.067 | 0.094 | 0.131 | 0.184
IRM | 0.045 | 0.052 | 0.122 | 0.161

###  2.2 Split Conformal Prediction Interval

In this section, we want to quantify the uncertainty of the estimates and
conduct a distribution-free prediction interval for given miscoverage rate
$\alpha$ using the invariant prediction and conformal inference. By solving
(2.1), we obtain the invariant estimates
$\hat{f}(x)=\hat{\omega}\cdot\hat{\Phi}(x)$ for any $x\in{\mathcal{X}}$.

We begin by applying the split conformal method to the data from different
environments, similar to method 1 in Dunn et al., (2022). Specifically, we
split the data in each training set into ${\mathcal{D}}_{tr}^{e}$ and
${\mathcal{D}}_{cal}^{e}$, where
${\mathcal{D}}_{cal}=\\{{\mathcal{D}}_{cal}^{e}\\}_{e=1}^{m}$ and
${\mathcal{D}}_{tr}=\\{{\mathcal{D}}_{tr}^{e}\\}_{e=1}^{m}$, with sample sizes
$n_{cal}^{e}$ and $n_{tr}^{e}$, respectively. We train the model (2.1) and
obtain $\hat{f}(x)$ using data in ${\mathcal{D}}^{e}_{tr}$ for $e=1,\ldots,m$.
The conformity score function is defined by $S(x,y)=|\hat{f}(x)-y|$, and we
use it to compute the conformity score
$\varepsilon_{i}^{e}=S(y_{i}^{e},x_{i}^{e})$ for $i\in{\mathcal{D}}_{cal}$ and
$e=1,\ldots,m$. We then rank all the
$\varepsilon_{i}^{e}=S(y_{i}^{e},x_{i}^{e})$ in ascending order and define the
$100(1-\alpha)\%$ empirical quantile of $\varepsilon_{i}^{e}$ as
$Q_{(1-\alpha)}(\varepsilon_{i}^{e},{\mathcal{D}}_{cal})$. Finally, we define
a prediction interval as

$\hat{C}_{1-\alpha}^{SC}(x)=\\{y:S(y,x)\leq
Q_{(1-\alpha)}(\varepsilon_{i}^{e},{\mathcal{D}}_{cal})\\}.$ (2.3)

In each environment, we assume that $S(X^{e},Y^{e})\sim{\mathbb{P}}_{S}^{e}$.
Since the joint distribution of ${\mathbb{P}}_{XY}^{e}$ is not identical
across environments, the distribution ${\mathbb{P}}_{S}^{e}$ also differs. As
shown in Dunn et al., (2022), when we pool data from different environments
and denote the joint distribution by
${\mathbb{P}}_{XY}^{1},\ldots,{\mathbb{P}}_{XY}^{m}\sim\Pi$, a new observation
$(x_{n+1},y_{n+1})$ collected from a new distribution $P_{XY}^{m+1}\sim\Pi$
satisfies $(x_{n+1},y_{n+1})\sim\widetilde{\Pi}$, where $\widetilde{\Pi}=\int
Pd\Pi(P)$. The constructed $\hat{C}^{SC}(x)$ defined in (2.3) satisfies the
marginal coverage for new observations.

###### Theorem 2.1

For any new observation $(x_{n+1},y_{n+1})$ from $\widetilde{\Pi}$, it holds
that

${\mathbb{P}}\left(y_{n+1}\in\hat{C}_{1-\alpha}^{SC}(x)\right)\geq 1-\alpha.$

The above result indicates that while $(x_{i}^{e},y_{i}^{e})$ are not
exchangeable across environments, they are independently sampled and not
identical. To ensure marginal coverage on the pooled data distributions, we
compute the average of the pooled joint distributions $(X,Y)$ and the
empirical quantile on the pooled residuals across environments.

#### Remark

In addition to the proposed approach, Dunn et al., (2022) propose
alternatives, such as subsampling approaches, to construct prediction sets for
observations from distinct distributions. These methods aim to recover the
population distribution by either pooling all subpopulations or subsampling
across the training groups. However, a limitation of the subsampling method is
that marginal coverage holds only when the number of environments
$m>1/\alpha-1$, which requires large group sizes to achieve a small
miscoverage rate. To address this limitation, we propose an adaptive conformal
method that does not require the number of groups and is adaptive to shifts
among the groups.

###  2.3 Calibrated Prediction Inverval

Now we construct a conformal interval adapt to the distribution shifts across
the environments. We first fit the IRM model (2.1) using the training data
${\mathcal{D}}_{tr}$ to obtain estimates of $\hat{\omega}$ and $\hat{\Phi}$.
Next, we use the conformity score function $S(x,y)=|\hat{f}(x)-y|$ to compute
the conformity score $\varepsilon_{i}^{e}=S(y_{i}^{e},x_{i}^{e})$ for
$i\in{\mathcal{D}}_{cal}$ and $e=1,\ldots,m$. In contrast to split conformal
method in Section 2.2 where all data was pooled together, we now weight the
conformity score for each environment $e$ separately, following the approach
in Guan, (2023). To construct a weighted conformity score, we assign weights
to new observations $(x_{n+1},y_{n+1})$ based on the data in the calibration
set ${\mathcal{D}}_{cal}$. Specifically, we compute the weight for any new
input $x_{n+1}$ by considering the first and second moment differences of the
samples, similar to the idea of moment matching in Arbel and Prünster, (2017).
We calculate the mean and standard deviation of each invariant representation
sample $\hat{\Phi}(x_{i})$ in the calibration set
${{\mathcal{D}}_{tr}^{e}}_{e=1}^{m}$, denoted by $\mu(\hat{\Phi}(x_{i}))$ and
$V(\hat{\Phi}(x_{i}))$, respectively, i.e.,

$\mu(\hat{\Phi}(x_{i}))=d^{-1}\sum_{j=1}^{d}\hat{\Phi}(x_{ij}),V({\hat{\Phi}(x_{i})})=\\{\hat{\Phi}(x_{ij})-\mu(\hat{\Phi}(x_{i}))\\}^{1/2},$
(2.4)

for $i\in\\{{\mathcal{D}}_{tr}^{e}\\}_{e=1}^{m}.$ Then we calculate the mean
and variance of each invariant representation sample $\hat{\Phi}(x_{i})$ in
the calibration set ${{\mathcal{D}}_{tr}^{e}}_{e=1}^{m}$ to obtain
$\mu^{e}(\hat{\Phi})=(n^{e}_{cal})^{-1}\sum_{i=1}^{n^{e}_{cal}}\mu({\hat{\Phi}(x_{i})})$
and
$V^{e}(\hat{\Phi})=(n^{e}_{cal})^{-1}\sum_{i=1}^{n^{e}_{cal}}V({\hat{\Phi}(x_{i})})$,
which represent the average mean and variation in environment $e$. For a new
input $x_{n+1}$, we define $\tau(x_{n+1},e)$ as the similarity between
$x_{n+1}$ and the samples in environment $e$ based on their sample mean and
variance differences. Specifically, $\tau(x_{n+1},e)$ is computed as

$\tau(x_{n+1},e)=\exp(-|V(\hat{\Phi}(x_{n+1}))-V^{e}(\hat{\Phi})|)\exp(-|\mu(\hat{\Phi}(x_{n+1}))-\mu^{e}(\hat{\Phi})|),$
(2.5)

We then assign weights $w_{x_{n+1},e}$ to sample $x_{n+1}$ with respect to
environment $e$, which are constructed using the similarity measure
$\tau(x_{n+1},e)$ and normalized by the sum of similarities across all
environments, i.e.,

$w_{x_{n+1},e}=\frac{\tau(x_{n+1},e)}{\sum_{e=1}^{m}\tau(x_{n+1},e)}.$ (2.6)

Let $\hat{Q}_{(1-\alpha)}^{e}(\varepsilon_{i}^{e},{\mathcal{D}}_{cal}^{e})$ be
the $100(1-\alpha)\%$ empirical quantile of
$\\{\varepsilon_{i}^{e}\\}_{i=1}^{n^{e}_{cal}}$ in the $e$ environment. Next,
we compute the weighted conformity score using the weights $w_{x_{n+1},e}$ and
the $100(1-\alpha)\%$ empirical quantile
$\hat{Q}{(1-\alpha)}^{e}(\varepsilon{i}^{e},{\mathcal{D}}_{cal}^{e})$ of
${\varepsilon_{i}^{e}}{i=1}^{n^{e}{cal}}$ in environment $e$. Specifically,
the weighted conformity score is defined by

$\tilde{Q}_{(1-\alpha)}(x_{n+1})=\sum_{e=1}^{m}\hat{Q}_{(1-\alpha)}^{e}(\varepsilon_{i}^{e},{\mathcal{D}}_{cal}^{e})\times
w(x_{n+1},e).$ (2.7)

This $\tilde{Q}_{(1-\alpha)}(x_{n+1})$ weights conformity score in each
environment according to the the sample $x_{n+1}$. Finally, the adaptive
conformal interval for invariant representation (ACIR) is constructed as

$\widehat{C}^{ACIR}_{1-\alpha}=\left[\hat{\omega}\cdot\hat{\Phi}(x_{n+1})-\tilde{Q}_{(1-\alpha)}(x_{n+1}),\hat{\omega}\cdot\hat{\Phi}(x_{n+1})+\tilde{Q}_{(1-\alpha)}(x_{n+1})\right],$
(2.8)

where $\tilde{Q}{(1-\alpha)}(x_{n+1})$ represents the data-driven weight
function of $x_{n+1}$ that measures the discrepancy between the true value of
$y_{n+1}$ and invariant prediction $\hat{\omega}\cdot\hat{\Phi}(x_{n+1})$.
Algorithm 1 shows the implementation of the ACIR construction.

Algorithm 1 The Adaptive Conformal Interval for Invariant Representation(ACIR)

Input: Data $(x_{i}^{e},y_{i}^{e})$ from $e=1,\ldots,m$ environments,
$x_{n+1}$, the IRM model $\mathcal{A}$, and $\alpha\in(0,1)$
Output: A prediction interval $\widehat{C}(x_{n+1})$ for unobserved $y_{n+1}$.

1:Randomly split the training data into 2 subsets
${\mathcal{D}}_{tr}^{e},{\mathcal{D}}_{cal}^{e}$ with sample size
$n_{tr}^{e},n_{cal}^{e}$ in each environment.

2:Train $\mathcal{A}$ on all samples in
$\\{{\mathcal{D}}_{tr}^{e}\\}_{e=1}^{m}$:
$\hat{\omega},\hat{\Phi}\leftarrow\mathcal{A}\left(\\{(x_{i},y_{i})\\}_{i\in\\{{\mathcal{D}}_{tr}^{e}\\}_{e=1}^{m}}\right)$.

3:Compute $\varepsilon_{i}^{e}=|y_{i}^{e}-\hat{\omega}\hat{\Phi}(x_{i}^{e})|$
for each $i\in{\mathcal{D}}_{tr}^{e}$ and $e\in\\{1,\ldots,m\\}$.

4:Compute
$\hat{Q}_{(1-\alpha)}^{e}(\varepsilon_{i}^{e},{\mathcal{D}}_{cal}^{e})$ as the
$\lceil(1-\alpha(1+n_{cal}^{e}))\rceil$-th largest value in
$\\{\varepsilon_{i}^{e}\\}_{i=1}^{n^{e}_{cal}}$ in the $e$ environment.

5:Calculate the mean and standard deviation of each sample point in every
environment and determine the environment-wise average
$\\{V^{e}(\hat{\Phi})\\}_{e=1}^{m}$ and $\\{\mu^{e}(\hat{\Phi})\\}_{e=1}^{m}$
using equation (2.4).

6:For a new input $x_{n+1}$, compute the weight $w_{x_{n+1},e}$ according to
equation (2.6).

7:Compute the weighted conformity score $\tilde{Q}_{(1-\alpha)}(x_{n+1})$
applying equation (2.7).

8:Use $\tilde{Q}_{(1-\alpha)}(x_{n+1})$ to construct the prediction set at
$x_{n+1}$ as:
$\widehat{C}^{ACIR}_{1-\alpha}=\left[\hat{\omega}\cdot\hat{\Phi}(x_{n+1})-\tilde{Q}_{(1-\alpha)}(x_{n+1}),\hat{\omega}\cdot\hat{\Phi}(x_{n+1})+\tilde{Q}_{(1-\alpha)}(x_{n+1})\right]$.

The proposed prediction interval $\widehat{C}^{ACIR}_{1-\alpha}(x)$ is
adaptive to changes in data distribution over time when operating with online
streaming data. When new data $x_{t+1}$ is received, we assign a weighted
conformity score to the observation and construct a conformal interval that
adapts to shifts among the environments.

###### Theorem 2.2

If $\|\hat{\Phi}-{\Phi}^{or}\|_{\infty}=o_{p}(1)$, where ${\Phi}^{or}$ is the
oracle $\Phi$ that satisfies invariance assumption, and new observation
$(x_{n+1}^{e},y_{n+1}^{e})$ is from $\widetilde{\Pi}$, it holds that

${\mathbb{P}}\left(y^{e}_{n+1}\in\widehat{C}_{1-\alpha}^{ACIR}(x^{e}_{n+1})\mid
x^{e}_{n+1}\right)=1-\alpha+o_{p}(1).$

Our results demonstrate that, under standard regularity conditions, the
conformal prediction set is nearly optimal and may provide asymptotic
conditional coverage, particularly when the initial model estimator is
accurate. This finding is consistent with previous studies Lei and Wasserman,
(2014) and Lei et al., (2018). However, in practice, the oracle representation
$\Phi$ is unknown. Therefore, we use the statistic ${\text{Inv}}(\Phi)$
defined by (2.2) to check the invariance assumption. Furthermore, we can
calibrate the invariant prediction by calculating the average distinction of
the invariance in the environment $e$ using $\delta^{e}=|(m-1)^{-1}\sum_{i\neq
e}\hat{M}_{e}(\Phi,X^{i})-\hat{M}_{e}(\Phi,X^{e})|$. When $\Phi=\Phi^{or}$,
${\mathbb{E}}(\delta^{e})=0$, and ${\mathbb{E}}(\delta^{e})$ increases as
$\|\Phi-\Phi^{or}\|_{\infty}$ increases. Therefore, we can add
$\sum_{e=1}^{m}w(x_{n+1},e)\delta^{e}$ when constructing the conformal
interval (2.8) to calibrate $\hat{\omega}\cdot\hat{\Phi}$ according to the
invariance. In our numerical experiments, we observed that the values of
$\delta^{e}$ were consistently small in relation to weighted conformity
scores. As a result, we decided to omit this step from our analysis presented
in Section 3.

##  3 Numerical results

In this section, we apply the proposed ACIR interval to the synthetic data and
a open-source financial data. We show the advantages of the proposed ACIR
compare to the split conformal without weight conformity score in both ERM and
IRM estimates.

###  3.1 Synthetic data

Here we present experiments on the (linear) structural equation model (SEM)
tasks introduced by Arjovsky et al., (2019) and Krueger et al., (2021). We
first simulate a synthetic dataset with a $10$-dimensional vector of $X$ that
contains causal effects $X_{1}$ and non-causal effect $X_{2}$, and a
continuous response $Y$ from the distributions as below.

$\displaystyle H^{e}$ $\displaystyle\sim{\mathcal{N}}(0,e^{2})$ $\displaystyle
X_{1}$ $\displaystyle\sim{\mathcal{N}}(0,e^{2})+W_{H\rightarrow 1}H^{e}$
$\displaystyle Y$ $\displaystyle\sim W_{1\rightarrow
Y}X_{1}+{\mathcal{N}}(0,\sigma_{y}^{2})+W_{H\rightarrow 1}H^{e}$
$\displaystyle X_{2}$ $\displaystyle\sim W_{Y\rightarrow
2}Y+{\mathcal{N}}(0,\sigma_{2}^{2})$

The variance of these distributions may vary across domains. We consider the
following two cases,

  1. 1.

Fully-observed (F), where $W_{h\rightarrow 1}=W_{h\rightarrow
y}=W_{h\rightarrow 2}=0$, or partially-observed (P), where
$\left(W_{h\rightarrow 1},W_{h\rightarrow y},W_{h\rightarrow 2}\right)$ are
Gaussian.

  2. 2.

Homoskedastic (O) $Y$-noise, where $\sigma_{y}^{2}=e^{2}$ and
$\sigma_{2}^{2}=1$, or heteroskedastic (E) $Y$-noise, where $\sigma_{y}^{2}=1$
and $\sigma_{2}^{2}=e^{2}$.

To simplify the analysis, we consider four settings denoted by 1-4: POU, PEU,
FOU, and FEU, where U represents the unscrambled $(X_{1},X_{2})$. Both $X_{1}$
and $X_{2}$ are generated as five-dimensional vectors and the environments
parameter $e$ is set to $[0.2,2.0,5.0]$.

Our method is applied to 6,000 independent observations from the described SEM
model, with 2,000 used to train the IRMv1 and ERM models, and another 2,000
for calibration. The remaining data is reserved for testing. We construct
split conformal intervals (SC) defined in (2.3) and adaptive conformal
intervals (AC) defined in (2.8) at a miscoverage rate of $\alpha=0.05$. The
performance of the prediction intervals is evaluated on the test set using the
following statistics: the average interval length is calculated by
${T}^{-1}\sum_{t=1}^{T}\tilde{Q}_{1-\alpha}(x_{t})$, and the coverage rate
calculated by
${T}^{-1}\sum_{t=1}^{T}\mathbbm{1}_{Y_{t}\in\hat{C}_{1-\alpha}\left(x_{t}\right)}$.
In addition, we perform 20 replications by randomly splitting the training
data 20 times and calculate the average statistics and their standard
deviations to draw a boxplot. The performance of prediction intervals based on
SC and AC under different settings is displayed in Figures 1 and 2.

Figure 1 shows the average performance on pooled data from all environments,
while Figure 2 displays the performance on each environment separately. In
Figure 1(a), the average length of the AC interval is consistently smaller
than that of the SC interval for both ERM and IRM estimates. Figure 1(b)
illustrates that the average coverage of the AC interval remains stable under
different settings, while the coverage in the SC method varies considerably
among different settings. This suggests that, on average, the AC method
achieves the desired coverage rate with smaller bandwidth under different
settings compared to the SC method.

Figure 2(a) reveals that the length of the AC interval is adaptive to the
environments, while the length of the SC interval is fixed, and the AC length
is always smaller than the SC length in every environment. In Figure 2(b), the
coverage of the SC interval performs worse than the AC interval and results in
an extreme coverage of $100\%$ in the heteroskedastic settings 2 and 4. While
the coverage of the AC interval varies in different environments for the IRM
estimates, it remains stable for ERM estimates. However, the coverage of the
AC method performs better than the SC method for both IRM and ERM estimates,
especially in the homoskedastic setting. This suggests that, in each
environment, the AC method achieves the desired coverage rate with a narrower
bandwidth under different settings compared to the SC method.

(a) Average length of different methods under different settings.

(b) Coverage of different methods under different settings.

Figure 1: The average length and coverage value of SC and AC method on ERM and
IRM models under different settings. Blue box is the SC interval for ERM
estimates, Blue box is the SC interval for ERM estimates,,Blue box is the SC
interval for ERM estimates, and Blue box is the SC interval for ERM estimates.

(a) Environment length of different methods under different settings.

(b) Environment coverage of different methods under different settings.

Figure 2: The average length and coverage value of SC and AC method in each
environment on ERM and IRM models under different settings. The labels are the
same as Figure 1.

###  3.2 Real data

This section applies the ACIR method to financial
data111https://www.kaggle.com/code/cnic92/explore-and-clean-financial-
indicators-dataset/notebook using the neural network setting in Krueger et
al., (2021) and cleaned data consisting of factors in the U.S. stock market
over five years, resulting in 37 features of company basic information and a
target variable representing the variation of stock prices. The training data
includes stock data from 2014 to 2016, while the testing data includes stock
data from 2017 to 2018, with each year treated as an environment and varying
sample sizes in each year.

We begin by splitting the training data in each environment into equal
calibration and proper training sets. We use the training set to fit the IRM
model to obtain the invariant predictions. Then, we use the proposed method to
construct adaptive conformal (AC) and split conformal (SC) prediction
intervals for the stock price variation in 2017 and 2018 at a miscoverage rate
of $\alpha=20\%$. The average length and coverage, as defined in Section 3.1,
are computed and presented in Table 3. Additionally, we calculate the coverage
and length in 2018 to compare the performance of the two methods in individual
environments. Table 3 shows that both AC and SC achieve marginal coverage,
while the AC method has a smaller length.

Table 3: The Performance of AC and SC method on financial data. Method | Coverage | Length | Coverage in 2018 | Length in 2018
---|---|---|---|---
SC | $80.7\%$ | 21.5 | $83.4\%$ | 21.5
AC | $79.5\%$ | 17.4 | $78.1\%$ | 15.4

##  4 Concluding remarks

This paper proposes an approach for uncertainty quantification of estimates
obtained through Invariant Risk Minimization (IRM). To identify the
environment to which a single test sample belongs, we propose a weight that
indicates the similarity between the sample and the environment’s sample
space. This weight can be combined with environment-wise statistics, including
environment-wise conformity scores, to obtain a weighted conformity score. We
construct an adaptive conformal interval using the weighted conformity score,
which is adaptive to distribution shifts across different environments. We
demonstrate that the adaptive conformal interval enjoys a conditional coverage
guarantee under conditions that can be checked by invariance assessment.

The proposed method and its theoretical properties have a main limitation in
that the weighted conformity score is defined by the first (mean) and second
(standard deviation) moment similarity of the samples. This approach may not
capture the characteristics of the sample distribution when using second
moments, such as in the case of skewed data. To better capture the skewness,
future research should consider alternative approaches that incorporate
higher-order moments and are more reliable and reasonable. Additionally, while
we propose a statistic for the assessment of invariance, it currently lacks
theoretical and numerical details. Future work should consider using this
assessment to calibrate the invariant representation and provide more thorough
theoretical and numerical analysis of its performance.

## References

  * Ahuja et al., (2021) Ahuja, K., Caballero, E., Zhang, D., Gagnon-Audet, J.-C., Bengio, Y., Mitliagkas, I., and Rish, I. (2021).  Invariance principle meets information bottleneck for out-of-distribution generalization.  Advances in Neural Information Processing Systems, 34:3438–3450.
  * Ahuja et al., (2020) Ahuja, K., Shanmugam, K., Varshney, K., and Dhurandhar, A. (2020).  Invariant risk minimization games.  In International Conference on Machine Learning, pages 145–155. PMLR.
  * Arbel and Prünster, (2017) Arbel, J. and Prünster, I. (2017).  A moment-matching ferguson & klass algorithm.  Statistics and Computing, 27:3–17.
  * Arjovsky et al., (2019) Arjovsky, M., Bottou, L., Gulrajani, I., and Lopez-Paz, D. (2019).  Invariant risk minimization.  arXiv preprint arXiv:1907.02893.
  * Barber et al., (2022) Barber, R. F., Candes, E. J., Ramdas, A., and Tibshirani, R. J. (2022).  Conformal prediction beyond exchangeability.  arXiv preprint arXiv:2202.13415.
  * Cauchois et al., (2020) Cauchois, M., Gupta, S., Ali, A., and Duchi, J. C. (2020).  Robust validation: Confident predictions even when distributions shift.  arXiv preprint arXiv:2008.04267.
  * Chang et al., (2020) Chang, S., Zhang, Y., Yu, M., and Jaakkola, T. (2020).  Invariant rationalization.  In International Conference on Machine Learning, pages 1448–1458. PMLR.
  * Chen et al., (2016) Chen, X., Monfort, M., Liu, A., and Ziebart, B. D. (2016).  Robust covariate shift regression.  In Artificial Intelligence and Statistics, pages 1270–1279. PMLR.
  * Creager et al., (2021) Creager, E., Jacobsen, J.-H., and Zemel, R. (2021).  Environment inference for invariant learning.  In International Conference on Machine Learning, pages 2189–2200. PMLR.
  * Dunn et al., (2022) Dunn, R., Wasserman, L., and Ramdas, A. (2022).  Distribution-free prediction sets for two-layer hierarchical models.  Journal of the American Statistical Association, pages 1–12.
  * Gibbs and Candes, (2021) Gibbs, I. and Candes, E. (2021).  Adaptive conformal inference under distribution shift.  Advances in Neural Information Processing Systems, 34:1660–1672.
  * Guan, (2023) Guan, L. (2023).  Localized conformal prediction: A generalized inference framework for conformal prediction.  Biometrika, 110(1):33–50.
  * Kivaranovic et al., (2020) Kivaranovic, D., Johnson, K. D., and Leeb, H. (2020).  Adaptive, distribution-free prediction intervals for deep networks.  In International Conference on Artificial Intelligence and Statistics, pages 4346–4356. PMLR.
  * Koyama and Yamaguchi, (2020) Koyama, M. and Yamaguchi, S. (2020).  Out-of-distribution generalization with maximal invariant predictor.  CoRR.
  * Krueger et al., (2021) Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Binas, J., Zhang, D., Le Priol, R., and Courville, A. (2021).  Out-of-distribution generalization via risk extrapolation (rex).  In International Conference on Machine Learning, pages 5815–5826. PMLR.
  * Lei et al., (2018) Lei, J., G’Sell, M., Rinaldo, A., Tibshirani, R. J., and Wasserman, L. (2018).  Distribution-free predictive inference for regression.  Journal of the American Statistical Association, 113(523):1094–1111.
  * Lei et al., (2013) Lei, J., Robins, J., and Wasserman, L. (2013).  Distribution-free prediction sets.  Journal of the American Statistical Association, 108(501):278–287.
  * Lei and Wasserman, (2014) Lei, J. and Wasserman, L. (2014).  Distribution-free prediction bands for non-parametric regression.  Journal of the Royal Statistical Society: Series B: Statistical Methodology, pages 71–96.
  * Li et al., (2022) Li, B., Shen, Y., Wang, Y., Zhu, W., Li, D., Keutzer, K., and Zhao, H. (2022).  Invariant information bottleneck for domain generalization.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 7399–7407.
  * Lin et al., (2022) Lin, Y., Dong, H., Wang, H., and Zhang, T. (2022).  Bayesian invariant risk minimization.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16021–16030.
  * Mahajan et al., (2021) Mahajan, D., Tople, S., and Sharma, A. (2021).  Domain generalization using causal matching.  In International Conference on Machine Learning, pages 7313–7324. PMLR.
  * Quinonero-Candela et al., (2008) Quinonero-Candela, J., Sugiyama, M., Schwaighofer, A., and Lawrence, N. D. (2008).  Dataset shift in machine learning.  Mit Press.
  * Romano et al., (2019) Romano, Y., Patterson, E., and Candes, E. (2019).  Conformalized quantile regression.  Advances in Neural Information Processing Systems, 32:3543–3553.
  * Rosenfeld et al., (2020) Rosenfeld, E., Ravikumar, P., and Risteski, A. (2020).  The risks of invariant risk minimization.  arXiv preprint arXiv:2010.05761.
  * Sugiyama et al., (2007) Sugiyama, M., Krauledat, M., and Müller, K.-R. (2007).  Covariate shift adaptation by importance weighted cross validation.  Journal of Machine Learning Research, 8(5).
  * Sugiyama and Müller, (2005) Sugiyama, M. and Müller, K.-R. (2005).  Input-dependent estimation of generalization error under covariate shift.
  * Taori et al., (2020) Taori, R., Dave, A., Shankar, V., Carlini, N., Recht, B., and Schmidt, L. (2020).  Measuring robustness to natural distribution shifts in image classification.  Advances in Neural Information Processing Systems, 33:18583–18599.
  * Tibshirani et al., (2019) Tibshirani, R. J., Foygel Barber, R., Candes, E., and Ramdas, A. (2019).  Conformal prediction under covariate shift.  Advances in neural information processing systems, 32.
  * Vovk, (2012) Vovk, V. (2012).  Conditional validity of inductive conformal predictors.  In Asian conference on machine learning, pages 475–490. PMLR.
  * Vovk et al., (2005) Vovk, V., Gammerman, A., and Shafer, G. (2005).  Algorithmic learning in a random world.  Springer Science & Business Media.
  * Wang et al., (2022) Wang, H., Si, H., Li, B., and Zhao, H. (2022).  Provable domain generalization via invariant-feature subspace recovery.  In International Conference on Machine Learning, pages 23018–23033. PMLR.
  * Zaffran et al., (2022) Zaffran, M., Féron, O., Goude, Y., Josse, J., and Dieuleveut, A. (2022).  Adaptive conformal predictions for time series.  In International Conference on Machine Learning, pages 25834–25866. PMLR.
  * Zeni et al., (2020) Zeni, G., Fontana, M., and Vantini, S. (2020).  Conformal prediction: a unified review of theory and new challenges.  arXiv preprint arXiv:2005.07972.
  * Zhou et al., (2022) Zhou, X., Lin, Y., Zhang, W., and Zhang, T. (2022).  Sparse invariant risk minimization.  In International Conference on Machine Learning, pages 27222–27244. PMLR.